66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
from mlflow.entities import Dataset
|
|
|
|
|
|
def _check(dataset, name, digest, source_type, source, schema=None, profile=None):
|
|
assert isinstance(dataset, Dataset)
|
|
assert dataset.name == name
|
|
assert dataset.digest == digest
|
|
assert dataset.source_type == source_type
|
|
assert dataset.source == source
|
|
assert dataset.schema == schema
|
|
assert dataset.profile == profile
|
|
|
|
|
|
def test_creation_and_hydration():
|
|
name = "my_name"
|
|
digest = "my_digest"
|
|
source_type = "my_source_type"
|
|
source = "my_source"
|
|
schema = "my_schema"
|
|
profile = "my_profile"
|
|
dataset = Dataset(name, digest, source_type, source, schema, profile)
|
|
_check(dataset, name, digest, source_type, source, schema, profile)
|
|
|
|
as_dict = {
|
|
"name": name,
|
|
"digest": digest,
|
|
"source_type": source_type,
|
|
"source": source,
|
|
"schema": schema,
|
|
"profile": profile,
|
|
}
|
|
assert dict(dataset) == as_dict
|
|
|
|
proto = dataset.to_proto()
|
|
dataset2 = Dataset.from_proto(proto)
|
|
_check(dataset2, name, digest, source_type, source, schema, profile)
|
|
|
|
dataset3 = Dataset.from_dictionary(as_dict)
|
|
_check(dataset3, name, digest, source_type, source, schema, profile)
|
|
|
|
|
|
def test_absent_fields():
|
|
name = "my_name"
|
|
digest = "my_digest"
|
|
source_type = "my_source_type"
|
|
source = "my_source"
|
|
dataset = Dataset(name, digest, source_type, source)
|
|
_check(dataset, name, digest, source_type, source)
|
|
|
|
as_dict = {
|
|
"name": name,
|
|
"digest": digest,
|
|
"source_type": source_type,
|
|
"source": source,
|
|
"profile": None,
|
|
"schema": None,
|
|
}
|
|
assert dict(dataset) == as_dict
|
|
|
|
proto = dataset.to_proto()
|
|
dataset2 = Dataset.from_proto(proto)
|
|
_check(dataset2, name, digest, source_type, source)
|
|
|
|
dataset3 = Dataset.from_dictionary(as_dict)
|
|
_check(dataset3, name, digest, source_type, source)
|