1  from mlflow.entities import Dataset, DatasetInput, InputTag
 2  
 3  
 4  def _check(dataset_input, tags, dataset):
 5      assert isinstance(dataset_input, DatasetInput)
 6      assert dataset_input.tags == tags
 7      assert dataset_input.dataset == dataset
 8  
 9  
10  def test_creation_and_hydration():
11      key = "my_key"
12      value = "my_value"
13      tags = [InputTag(key, value)]
14      name = "my_name"
15      digest = "my_digest"
16      source_type = "my_source_type"
17      source = "my_source"
18      schema = "my_schema"
19      profile = "my_profile"
20      dataset = Dataset(name, digest, source_type, source, schema, profile)
21      dataset_input = DatasetInput(dataset=dataset, tags=tags)
22      _check(dataset_input, tags, dataset)
23  
24      as_dict = {"dataset": dataset, "tags": tags}
25      assert dict(dataset_input) == as_dict
26  
27      proto = dataset_input.to_proto()
28      dataset_input2 = DatasetInput.from_proto(proto)
29      _check(dataset_input2, tags, dataset)
30  
31      dataset_input3 = DatasetInput.from_dictionary(as_dict)
32      _check(dataset_input3, tags, dataset)