From bde614e2a68d6f7c7e20b9f0e71226f82d9f7692 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 4 Sep 2026 10:54:23 -0300 Subject: [PATCH] QTZPOC-20 QTZPOC-20: Refactor model repository tests for improved assertion logic and exception handling, including correcting pytest.raises exception access and refining mock verifications. --- .../utils/repository/test_model_repository.py | 129 ++++++++++-------- 1 file changed, 69 insertions(+), 60 deletions(-) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 38fab72..95c911c 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -105,11 +105,10 @@ def test_get_model_run_id_not_registered_models(mlflow_repository): with pytest.raises(mlflow_lib.exceptions.MlflowException) as e: mlflow_repository.get_model_run_id('test') - mlflow_repository.client.search_registered_models.assert_called_once_with( - filter_string="name='test'" - ) - - assert str(e.value) == "Model 'test' not found in the Model Registry." + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + assert str(e.value) == "Model 'test' not found in the Model Registry." def test_get_model_run_id_not_stage_versions(mlflow_repository): @@ -124,14 +123,13 @@ def test_get_model_run_id_not_stage_versions(mlflow_repository): with pytest.raises(mlflow_lib.exceptions.MlflowException) as e: mlflow_repository.get_model_run_id('test') - mlflow_repository.client.search_registered_models.assert_called_once_with( - filter_string="name='test'" - ) - mlflow_repository.client.search_model_versions.assert_called_once_with( - filter_string="name='test'" - ) - - assert str(e.value) == "Model 'test' in stage 'Production' not found in the Model Registry." + mlflow_repository.client.search_registered_models.assert_called_once_with( + filter_string="name='test'" + ) + mlflow_repository.client.search_model_versions.assert_called_once_with( + filter_string="name='test'" + ) + assert str(e.value) == "Model 'test' in stage 'Production' not found in the Model Registry." def test_get_model_run_id_success(mlflow_repository): @@ -231,7 +229,7 @@ def test_get_experiment_none_not_create(mlflow, mlflow_repository): with pytest.raises(ValueError) as e: mlflow_repository.get_experiment('test', create_if_not_exists=False) - assert str(e) == 'Experiment test not found' + assert str(e.value) == 'Experiment test not found' def test_get_model_params(mlflow, mlflow_repository): @@ -354,10 +352,10 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor with pytest.raises(ValueError): await mlflow_repository.dowload_artifacts('test', {}, 'path') - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + mlflow_repository.emit_metric.assert_called_once_with( + metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY + ) + mlflow_repository.observe_lag.assert_not_called() @pytest.mark.asyncio @@ -442,7 +440,7 @@ def test_get_experiment_create_error(mlflow, mlflow_repository): with pytest.raises(ValueError) as e: mlflow_repository.get_experiment('test', create_if_not_exists=True) - assert str(e) == 'Experiment test not found after creation, unknown reason' + assert str(e.value) == 'Experiment test not found after creation, unknown reason' @pytest.mark.asyncio @@ -475,12 +473,10 @@ async def test_load_predict_model_pyfunc(mlflow, mlflow_repository): async def test_load_predict_model_error(mlflow_repository): with pytest.raises(ValueError) as e: await mlflow_repository.load_predict_model('test_model', {}, 'invalid') - assert str(e) == INVALID_FLAVOR_MESSAGE - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + assert str(e.value) == INVALID_FLAVOR_MESSAGE + mlflow_repository.emit_metric.assert_not_called() + mlflow_repository.observe_lag.assert_not_called() TRANSFORM_MODEL_URI = 'runs:/run-123/data_model' @@ -540,24 +536,20 @@ async def test_load_transform_model_error(mlflow_repository): with pytest.raises(ValueError) as e: await mlflow_repository.load_transform_model('test_model', {}, 'invalid') - assert str(e) == INVALID_FLAVOR_MESSAGE - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + assert str(e.value) == INVALID_FLAVOR_MESSAGE + mlflow_repository.emit_metric.assert_not_called() + mlflow_repository.observe_lag.assert_not_called() @pytest.mark.asyncio async def test_download_model_invalid_model_type(mlflow_repository): with pytest.raises(ValueError) as e: await mlflow_repository.download_model('test_model', {}, 'invalid', 'sklearn') - assert str(e) == "Invalid model_type. Use 'predict' or 'transform'." - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + assert str(e.value) == "Invalid model_type. Use 'predict' or 'transform'." + mlflow_repository.emit_metric.assert_not_called() + mlflow_repository.observe_lag.assert_not_called() @pytest.mark.asyncio @@ -618,11 +610,18 @@ async def test_download_model_prepares_the_code_dir_first( mlflow_repository, model_type, artifact_name ): """The `code` step runs before any loader, for the artifact matching the model type.""" - calls = [] - mlflow_repository.download_artifact_code_dir = AsyncMock( - side_effect=lambda *_a, **_k: calls.append('code') or 'code-artifact-dir' - ) - loader = AsyncMock(side_effect=lambda *_a: calls.append('load') or ('model', None)) + calls: list[str] = [] + + async def record_code_dir(*_a, **_k): + calls.append('code') + return 'code-artifact-dir' + + async def record_load(*_a): + calls.append('load') + return 'model', None + + mlflow_repository.download_artifact_code_dir = AsyncMock(side_effect=record_code_dir) + loader = AsyncMock(side_effect=record_load) mlflow_repository.load_predict_model = loader mlflow_repository.load_transform_model = loader @@ -890,7 +889,8 @@ async def test_get_cached_operation_invalid_operation(force_memory_release, mlfl await mlflow_repository.get_cached_operation( 'model_name', data, 'invalid', 0, 'sklearn', {} ) - assert str(e) == "Invalid operation. Use 'transform' or 'predict'." + + assert str(e.value) == "Invalid operation. Use 'transform' or 'predict'." @patch('laborious.utils.repository.model_repository.pd.merge') @@ -1149,11 +1149,12 @@ async def test_log_model_error(mlflow_repository): await mlflow_repository.log_model( model_data, 'invalid', 'prediction_model', metadata['metadata'] ) - assert str(e) == INVALID_FLAVOR_MESSAGE - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + + assert str(e.value) == INVALID_FLAVOR_MESSAGE + mlflow_repository.emit_metric.assert_called_once_with( + metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY + ) + mlflow_repository.observe_lag.assert_not_called() @patch('laborious.utils.repository.model_repository.force_memory_release') @@ -1302,10 +1303,10 @@ async def test_create_new_experiment_error( 'pyfunc', ) - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + mlflow_repository.emit_metric.assert_called_once_with( + metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY + ) + mlflow_repository.observe_lag.assert_not_called() @pytest.mark.asyncio @@ -1360,24 +1361,32 @@ async def test_update_production_model_by_run_id_error_register_model(mlflow, ml with pytest.raises(ValueError): await mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + mlflow_repository.emit_metric.assert_called_once_with( + metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY + ) + mlflow_repository.observe_lag.assert_not_called() @pytest.mark.asyncio async def test_update_production_model_by_run_id_error_transition_model_version_stage( mlflow, mlflow_repository ): + mlflow_repository.client.get_registered_model.return_value = MagicMock( + latest_versions=[MagicMock(version='1')] + ) mlflow_repository.client.transition_model_version_stage.side_effect = ValueError('error') with pytest.raises(ValueError): await mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) - mlflow_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY - ) - mlflow_repository.observe_lag.assert_not_called() + # register_model succeeded before transition failed, so its WRITE_COUNT/WRITE_LAG + # were already recorded; only the transition step reports as an error. + mlflow_repository.emit_metric.assert_has_calls( + [ + call(metric_object=metrics.MODEL_WRITE_COUNT, tags=ANY), + call(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY), + ] + ) + mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY) @pytest.mark.asyncio @@ -2131,7 +2140,7 @@ def test_load_model_pickle_adds_code_dir_before_load( ): code_dir = build_code_dir(tmp_path) local_dir = build_artifact_dir(tmp_path, 'model.pkl') - sys_path_at_load = [] + sys_path_at_load: list[str] = [] joblib_mock.load.side_effect = lambda _path: sys_path_at_load.extend(sys.path) mlflow_repository._load_model_pickle(local_dir, 'prediction_model') @@ -2286,7 +2295,7 @@ async def test_download_model_puts_code_dir_on_sys_path_before_the_loader( """MLflow extends `sys.path` only when `MLmodel` declares `code`; the directory decides here.""" code_dir = stub_remote_code_dir(mlflow_repository, tmp_path, 'prediction_model') mlflow_repository.get_model_uri = MagicMock(return_value=TRANSFORM_MODEL_URI) - sys_path_at_load = [] + sys_path_at_load: list[str] = [] getattr(mlflow, flavor).load_model.side_effect = observe_sys_path(sys_path_at_load) _model, artifact_dir = await mlflow_repository.download_model( @@ -2339,7 +2348,7 @@ async def test_download_model_load_wrapper_puts_code_dir_on_sys_path( """The wrapper branch loads a downloaded artifact, so `code` has to be in place first too.""" code_dir = stub_remote_code_dir(mlflow_repository, tmp_path, 'prediction_model') mlflow_repository.dowload_artifacts = AsyncMock(return_value=str(tmp_path)) - sys_path_at_load = [] + sys_path_at_load: list[str] = [] mlflow.pyfunc.load_model.side_effect = observe_sys_path(sys_path_at_load) await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True)