def test_update_model_version(mock_get_request_message, mock_model_registry_store):
rm = RegisteredModel("model1")
mv = ModelVersion(registered_model=rm, version=32)
mock_get_request_message.return_value = UpdateModelVersion(model_version=mv.to_proto(),
stage="Production",
description="Great model!")
_update_model_version()
args, _ = mock_model_registry_store.update_model_version.call_args
assert args == (mv, "Production", "Great model!")
def test_update_model_version(mock_get_request_message, mock_model_registry_store):
name = "model1"
version = "32"
description = "Great model!"
mock_get_request_message.return_value = UpdateModelVersion(name=name, version=version,
description=description)
mv = ModelVersion(name=name, version=version, creation_timestamp=123, description=description)
mock_model_registry_store.update_model_version.return_value = mv
_update_model_version()
_, args = mock_model_registry_store.update_model_version.call_args
assert args == {"name": name, "version": version, "description": description}