1  from mlflow.entities import RunData
 2  
 3  
 4  def _check_metrics(metric_objs, metrics_dict, expected_metrics):
 5      assert {m.key for m in metric_objs} == {m.key for m in expected_metrics}
 6      assert {m.value for m in metric_objs} == {m.value for m in expected_metrics}
 7      assert {m.timestamp for m in metric_objs} == {m.timestamp for m in expected_metrics}
 8      assert {m.step for m in metric_objs} == {m.step for m in expected_metrics}
 9      assert len(metrics_dict) == len(expected_metrics)
10      assert metrics_dict == {m.key: m.value for m in expected_metrics}
11  
12  
13  def _check_params(params_dict, expected_params):
14      assert params_dict == {p.key: p.value for p in expected_params}
15  
16  
17  def _check_tags(tags_dict, expected_tags):
18      assert tags_dict == {t.key: t.value for t in expected_tags}
19  
20  
21  def _check(rd, metrics, params, tags):
22      assert isinstance(rd, RunData)
23      _check_metrics(rd._metric_objs, rd.metrics, metrics)
24      _check_params(rd.params, params)
25      _check_tags(rd.tags, tags)
26  
27  
28  def test_creation_and_hydration(run_data):
29      rd, metrics, params, tags = run_data
30      _check(rd, metrics, params, tags)
31      as_dict = {
32          "metrics": {m.key: m.value for m in metrics},
33          "params": {p.key: p.value for p in params},
34          "tags": {t.key: t.value for t in tags},
35      }
36      assert dict(rd) == as_dict
37      proto = rd.to_proto()
38      rd2 = RunData.from_proto(proto)
39      _check(rd2, metrics, params, tags)