def test_update_run_info_invalid_run_id(mocker):
mock_client = mocker.Mock()
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(ValueError):
store.update_run_info("invalid-run-id", RunStatus.RUNNING,
RUN_ENDED_AT_MILLISECONDS)
def test_update_run_info_client_error(mocker):
mock_client = mocker.Mock()
mock_client.update_run_info.side_effect = HttpError(
mocker.Mock(), "Experiment run with ID _ not found in project _")
mocker.patch("faculty.client", return_value=mock_client)
store = FacultyRestStore(STORE_URI)
with pytest.raises(
MlflowException,
match="Experiment run with ID _ not found in project _",
):
store.update_run_info(RUN_UUID_HEX_STR, "RUNNING",
RUN_ENDED_AT_MILLISECONDS)
def test_update_run_info(mocker, run_status, faculty_run_status):
faculty_run = mocker.Mock()
mock_client = mocker.Mock()
mock_client.update_run_info.return_value = faculty_run
mocker.patch("faculty.client", return_value=mock_client)
mlflow_run_info = mocker.Mock()
mlflow_run = mocker.Mock()
mlflow_run.info = mlflow_run_info
run_converter_mock = mocker.patch(
"mlflow_faculty.tracking.faculty_run_to_mlflow_run",
return_value=mlflow_run,
)
store = FacultyRestStore(STORE_URI)
returned_run_info = store.update_run_info(
RUN_UUID_HEX_STR, run_status, RUN_ENDED_AT_MILLISECONDS
)
mock_client.update_run_info.assert_called_once_with(
PROJECT_ID, RUN_UUID, faculty_run_status, RUN_ENDED_AT
)
run_converter_mock.assert_called_once_with(faculty_run)
assert returned_run_info == mlflow_run_info
