1  from mlflow.entities import RunInputs
 2  from mlflow.entities.dataset_input import DatasetInput
 3  
 4  
 5  def _check_inputs(run_datasets, datasets):
 6      for d1, d2 in zip(run_datasets, datasets):
 7          assert d1.dataset.digest == d2.dataset.digest
 8          assert d1.dataset.name == d2.dataset.name
 9          assert d1.dataset.source_type == d2.dataset.source_type
10          assert d1.dataset.source == d2.dataset.source
11          for t1, t2 in zip(d1.tags, d2.tags):
12              assert t1.key == t2.key
13              assert t1.value == t2.value
14  
15  
16  def _check(inputs, datasets):
17      assert isinstance(inputs, RunInputs)
18      _check_inputs(inputs.dataset_inputs, datasets)
19  
20  
21  def test_creation_and_hydration(run_inputs):
22      run_inputs, datasets = run_inputs
23      _check(run_inputs, datasets)
24      as_dict = {
25          "dataset_inputs": [
26              {
27                  "dataset": {
28                      "digest": "digest1",
29                      "name": "name1",
30                      "profile": None,
31                      "schema": None,
32                      "source": "source",
33                      "source_type": "my_source_type",
34                  },
35                  "tags": {"key": "value"},
36              }
37          ],
38          "model_inputs": [],
39      }
40      assert run_inputs.to_dictionary() == as_dict
41      proto = run_inputs.to_proto()
42      run_inputs2 = RunInputs.from_proto(proto)
43      _check(run_inputs2, datasets)
44      assert isinstance(run_inputs2.dataset_inputs[0], DatasetInput)