/ tests / utils / test_unity_catalog_oss_utils.py
test_unity_catalog_oss_utils.py
 1  from mlflow.entities.model_registry import RegisteredModel
 2  from mlflow.protos.unity_catalog_oss_messages_pb2 import (
 3      RegisteredModelInfo,
 4  )
 5  from mlflow.utils._unity_catalog_oss_utils import get_registered_model_from_uc_oss_proto
 6  
 7  
 8  def test_registered_model_from_uc_oss_proto():
 9      expected_registered_model = RegisteredModel(
10          name="catalog.schema.name",
11          creation_timestamp=1,
12          last_updated_timestamp=2,
13          description="description",
14      )
15  
16      uc_oss_proto = RegisteredModelInfo(
17          name="name",
18          catalog_name="catalog",
19          schema_name="schema",
20          created_at=1,
21          updated_at=2,
22          comment="description",
23      )
24  
25      actual_registered_model = get_registered_model_from_uc_oss_proto(uc_oss_proto)
26      assert actual_registered_model == expected_registered_model