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.

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(
+    mlflow_repository.client.search_registered_models.assert_called_once_with(
-            filter_string="name='test'"
+        filter_string="name='test'"
-        )
+    )
-
+    assert str(e.value) == "Model 'test' not found in the Model Registry."
-        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(
+    mlflow_repository.client.search_registered_models.assert_called_once_with(
-            filter_string="name='test'"
+        filter_string="name='test'"
-        )
+    )
-        mlflow_repository.client.search_model_versions.assert_called_once_with(
+    mlflow_repository.client.search_model_versions.assert_called_once_with(
-            filter_string="name='test'"
+        filter_string="name='test'"
-        )
+    )
-
+    assert str(e.value) == "Model 'test' in stage 'Production' not found in the Model Registry."
-        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(
+    mlflow_repository.emit_metric.assert_called_once_with(
-            metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
+        metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
-        )
+    )
-        mlflow_repository.observe_lag.assert_not_called()
+    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(
+    assert str(e.value) == INVALID_FLAVOR_MESSAGE
-            metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
+    mlflow_repository.emit_metric.assert_not_called()
-        )
+    mlflow_repository.observe_lag.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(
+    assert str(e.value) == INVALID_FLAVOR_MESSAGE
-            metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
+    mlflow_repository.emit_metric.assert_not_called()
-        )
+    mlflow_repository.observe_lag.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(
+    assert str(e.value) == "Invalid model_type. Use 'predict' or 'transform'."
-            metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
+    mlflow_repository.emit_metric.assert_not_called()
-        )
+    mlflow_repository.observe_lag.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 = []
+    calls: list[str] = []
-    mlflow_repository.download_artifact_code_dir = AsyncMock(
+
-        side_effect=lambda *_a, **_k: calls.append('code') or 'code-artifact-dir'
+    async def record_code_dir(*_a, **_k):
-    )
+        calls.append('code')
-    loader = AsyncMock(side_effect=lambda *_a: calls.append('load') or ('model', None))
+        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(
+    assert str(e.value) == INVALID_FLAVOR_MESSAGE
-            metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
+    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.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(
+    mlflow_repository.emit_metric.assert_called_once_with(
-            metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
+        metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
-        )
+    )
-        mlflow_repository.observe_lag.assert_not_called()
+    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(
+    mlflow_repository.emit_metric.assert_called_once_with(
-            metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
+        metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
-        )
+    )
-        mlflow_repository.observe_lag.assert_not_called()
+    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(
+    # register_model succeeded before transition failed, so its WRITE_COUNT/WRITE_LAG
-            metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY
+    # were already recorded; only the transition step reports as an error.
-        )
+    mlflow_repository.emit_metric.assert_has_calls(
-        mlflow_repository.observe_lag.assert_not_called()
+        [
+            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)