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