示例#1
0
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}
示例#2
0
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, )