Python _update_model_versionの例

コード例 #1

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!")

コード例 #2

ファイル: test_handlers.py プロジェクト: zirubak/mlflow

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}