Python _update_registered_model Examples

Example #1

File: test_handlers.py Project: zirubak/mlflow

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)}

Example #2

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)}