def test_update_registered_model(mock_get_request_message, mock_model_registry_store):
name = "model_1"
description = "Test model"
mock_get_request_message.return_value = UpdateRegisteredModel(name=name,
description=description)
rm2 = RegisteredModel(name, description=description)
mock_model_registry_store.update_registered_model.return_value = rm2
resp = _update_registered_model()
_, args = mock_model_registry_store.update_registered_model.call_args
assert args == {"name": name, "description": u"Test model"}
assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm2)}
def test_update_registered_model(mock_get_request_message, mock_model_registry_store):
rm1 = RegisteredModel("model_1")
mock_get_request_message.return_value = UpdateRegisteredModel(registered_model=rm1.to_proto(),
name="model_2",
description="Test model")
rm2 = RegisteredModel("model_2")
mock_model_registry_store.update_registered_model.return_value = rm2
resp = _update_registered_model()
args, _ = mock_model_registry_store.update_registered_model.call_args
assert args == (rm1, u"model_2", u"Test model")
assert json.loads(resp.get_data()) == {"registered_model": jsonify(rm2)}