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

SIENTIAPDE-1273

Browse Source 
SIENTIAPDE-1273
Enhance security analysis and SQL injection handling

- Added skip for potential SQL injection false positives in Bandit configuration.
- Updated validate.sh to use the pyproject.toml configuration for Bandit security analysis.
- Refactored code to replace ensure_dataframe utility with direct DataFrame usage in multiple activities, improving clarity and reducing dependencies.
- Removed the deprecated dataframe_utils module to streamline the codebase.
This commit is contained in:
Vitor Pimentel dos Santos
2025-11-17 16:04:54 -03:00
parent 1014c33dd9
commit a88a15c60a
18 changed files with 516 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -193,9 +193,11 @@ def test_get_model_params(mlflow, mlflow_repository):
def test_check_artifact_exists_true(mlflow_repository):
artifact = MagicMock(path='test_artifact')
mlflow_repository.client.list_artifacts.return_value = [artifact]
result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
result = mlflow_repository.check_artifact_exists(
'run_id', 'test_artifact', metadata['metadata']
)
assert result is True
mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
@@ -203,9 +205,11 @@ def test_check_artifact_exists_true(mlflow_repository):
def test_check_artifact_exists_false(mlflow_repository):
artifact = MagicMock(path='other_artifact')
mlflow_repository.client.list_artifacts.return_value = [artifact]
result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
result = mlflow_repository.check_artifact_exists(
'run_id', 'test_artifact', metadata['metadata']
)
assert result is False
mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
@@ -299,16 +303,22 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor
@patch('laborious.utils.repository.model_repository.mlflow')
@patch('laborious.utils.repository.model_repository.pd')
@patch('laborious.utils.repository.model_repository.StringIO')
async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repository):
async def test_load_artifact_dataframe_success(_stringio, pd, mlflow, mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
mlflow.artifacts.load_text.return_value = 'col1,col2\n1,2\n3,4'
result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
result = await mlflow_repository.load_artifact_dataframe(
'model_name', 'artifact_path', metadata['metadata']
)
mlflow_repository.get_model_run_id.assert_called_once_with(
model_name='model_name', stage='Production'
)
mlflow_repository.check_artifact_exists.assert_called_once_with(
'run_id', 'artifact_path', metadata['metadata']
)
mlflow.artifacts.load_text.assert_called_once_with('runs:/run_id/artifact_path')
assert result == pd.read_csv.return_value
mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY)
@@ -321,12 +331,18 @@ async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repo
async def test_load_artifact_dataframe_not_exists(mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
mlflow_repository.check_artifact_exists = MagicMock(return_value=False)
result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
result = await mlflow_repository.load_artifact_dataframe(
'model_name', 'artifact_path', metadata['metadata']
)
assert result is None
mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
mlflow_repository.get_model_run_id.assert_called_once_with(
model_name='model_name', stage='Production'
)
mlflow_repository.check_artifact_exists.assert_called_once_with(
'run_id', 'artifact_path', metadata['metadata']
)
@pytest.mark.asyncio
@@ -335,10 +351,12 @@ async def test_load_artifact_dataframe_error(mlflow, mlflow_repository):
mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
mlflow.artifacts.load_text.side_effect = ValueError('error')
with pytest.raises(ValueError):
await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
await mlflow_repository.load_artifact_dataframe(
'model_name', 'artifact_path', metadata['metadata']
)
mlflow_repository.emit_metric.assert_called_once_with(
metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
)
@@ -1059,7 +1077,9 @@ async def test_create_new_experiment(
data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=False)
# Verify prediction_data.to_csv was called with correct arguments
prediction_data.to_csv.assert_called_once()
assert prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
assert (
prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
)
assert prediction_data.to_csv.call_args[1]['index'] is False
mlflow.start_run.assert_called_once_with(
@@ -1089,10 +1109,12 @@ async def test_create_new_experiment(
}
)
mlflow.log_artifact.assert_has_calls([
call('./tmp/artifacts/model_name/retrain_data.csv'),
call('./tmp/artifacts/model_name/evaluation_data.csv'),
])
mlflow.log_artifact.assert_has_calls(
[
call('./tmp/artifacts/model_name/retrain_data.csv'),
call('./tmp/artifacts/model_name/evaluation_data.csv'),
]
)
force_memory_release.assert_called_once_with(mlflow_repository.logger)
@@ -1472,9 +1494,9 @@ def test_get_prediction_data_dataframe(mlflow_repository):
retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
target_name = 'target'
result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
prediction_model.predict.assert_called_once_with(retrain_dataset)
assert 'prediction' in result.columns
assert 'target' in result.columns
@@ -1487,9 +1509,9 @@ def test_get_prediction_data_array(mlflow_repository):
retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
prediction_model.predict.return_value = [5, 6]
target_name = 'target'
result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
prediction_model.predict.assert_called_once_with(retrain_dataset)
assert 'prediction' in result.columns
assert 'target' in result.columns