def test_delete_model_version(mock_get_request_message, mock_model_registry_store): name = "model1" version = "32" mock_get_request_message.return_value = DeleteModelVersion(name=name, version=version) _delete_model_version() _, args = mock_model_registry_store.delete_model_version.call_args assert args == {"name": name, "version": version}
def test_delete_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 = DeleteModelVersion(model_version=mv.to_proto()) _delete_model_version() args, _ = mock_model_registry_store.delete_model_version.call_args assert args == (mv, )