gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

QTZPOC-20

Browse Source 
QTZPOC-20: Refactor model repository tests for improved assertion logic and exception handling, including correcting pytest.raises exception access and refining mock verifications.
This commit is contained in:
Vitor Pimentel dos Santos
2026-09-04 10:54:23 -03:00
parent 6a412147ee
commit bde614e2a6
1 changed files with 69 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

129
tests/laborious/utils/repository/test_model_repository.py
View File

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