gitea_admin/sientia-dataops-laborious_temporal
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-aignosi
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

538
tests/laborious/activities/test_model_metrics.py
View File

@@ -1,4 +1,4 @@
from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
from unittest.mock import ANY, AsyncMock, MagicMock, patch
from pandas import DataFrame
from pytest import fixture, mark
@@ -22,7 +22,13 @@ def model_metrics_activity():
model_metrics.send_notification = MagicMock()
model_metrics.send_notification_async = AsyncMock()
model_metrics.emit_metric = AsyncMock()
model_metrics.get_core_labels = MagicMock(return_value={'pod_id': 'test_pod', 'model_name': 'test_model', 'workflow_name': 'test_workflow'})
model_metrics.get_core_labels = MagicMock(
return_value={
'pod_id': 'test_pod',
'model_name': 'test_model',
'workflow_name': 'test_workflow',
}
)
model_metrics.observe_lag = AsyncMock()
model_metrics.pod_id = 'test_pod'
return model_metrics
@@ -60,7 +66,7 @@ async def test_calculate_drift_invalid_chunk_period(model_metrics_activity):
try:
await model_metrics_activity.calculate_drift(input_data)
except ValueError as e:
assert str(e) == "Invalid chunk period: invalid, must be \"min\" or \"s\""
assert str(e) == 'Invalid chunk period: invalid, must be "min" or "s"'
model_metrics_activity.error.assert_called_once_with(
'Invalid chunk period: invalid', metadata['metadata']
)
@@ -76,26 +82,41 @@ async def test_calculate_drift_with_reference_data(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -124,16 +145,20 @@ async def test_calculate_drift_with_reference_data(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert result == mock_drift_df.to_dict.return_value
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@@ -144,18 +169,31 @@ async def test_calculate_drift_without_reference_data(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [False]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': False,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
target_data_dict = {
@@ -163,19 +201,25 @@ async def test_calculate_drift_without_reference_data(
'variable': ['feature1', 'feature1', 'feature1'],
'value': [1.0, 2.0, 3.0],
}
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
mock_target_df.reset_index.return_value = mock_target_df
mock_target_df.dropna.return_value = mock_target_df
mock_target_df.sort_values.return_value = mock_target_df
mock_target_df.head.return_value = DataFrame({'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]})
mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']
mock_target_df.head.return_value = DataFrame(
{'timestamp': ['2023-05-26 11:12:27'], 'feature1': [1.0]}
)
mock_target_df.__getitem__.return_value.apply.return_value = [
'2023-05-26 11:12:27',
'2023-05-26 11:12:28',
'2023-05-26 11:12:29',
]
mock_target_df.drop.return_value.columns = ['feature1']
mock_dataframe.return_value = mock_target_df
mock_dataframe.side_effect = lambda x=None: mock_target_df if x is not None else mock_target_df
input_data = {
**metadata,
'model_name': 'test_model',
@@ -191,7 +235,7 @@ async def test_calculate_drift_without_reference_data(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value
model_metrics_activity.warning.assert_called()
model_metrics_activity.send_notification_async.assert_called_once_with(
@@ -203,11 +247,15 @@ async def test_calculate_drift_without_reference_data(
attachment_content=ANY,
)
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@@ -218,15 +266,17 @@ async def test_calculate_drift_empty_drift_df(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=DataFrame())
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -235,7 +285,7 @@ async def test_calculate_drift_empty_drift_df(
mock_target_df.__getitem__.return_value.apply.return_value = ['2023-05-26 11:12:27']
mock_target_df.drop.return_value.columns = ['feature1']
mock_dataframe.return_value = mock_target_df
input_data = {
**metadata,
'model_name': 'test_model',
@@ -255,8 +305,10 @@ async def test_calculate_drift_empty_drift_df(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert result == {}
model_metrics_activity.warning.assert_called_with('No drift metrics found', metadata['metadata'])
assert result == []
model_metrics_activity.warning.assert_called_with(
'No drift metrics found', metadata['metadata']
)
@mark.asyncio
@@ -267,35 +319,37 @@ async def test_calculate_drift_empty_after_timestamp_filter(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
# Set up __getitem__ to handle filtering - timestamp access returns series with isin=False
# and filtering returns empty DataFrame
mock_timestamp_series = MagicMock()
mock_timestamp_series.isin.return_value = [False]
mock_empty_df = MagicMock()
mock_empty_df.empty = True
def getitem_side_effect(key):
if key == 'timestamp':
return mock_timestamp_series
else:
# This is the filtering operation - return empty DataFrame
return mock_empty_df
mock_drift_df.__getitem__.side_effect = getitem_side_effect
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -324,13 +378,13 @@ async def test_calculate_drift_empty_after_timestamp_filter(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert result == {}
assert result == []
model_metrics_activity.warning.assert_called_with(
'No drift metrics found after dropping rows where timestamp is not in target data',
metadata['metadata']
metadata['metadata'],
)
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
@@ -342,26 +396,41 @@ async def test_calculate_drift_success_min(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -390,46 +459,63 @@ async def test_calculate_drift_success_min(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert result == mock_drift_df.to_dict.return_value
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@patch('laborious.activities.model_metrics.DataFrame')
@patch('laborious.activities.model_metrics.to_datetime')
async def test_calculate_drift_success_s(
mock_to_datetime, mock_dataframe, model_metrics_activity
):
async def test_calculate_drift_success_s(mock_to_datetime, mock_dataframe, model_metrics_activity):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_to_datetime.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df = MagicMock()
mock_drift_df.empty = False
mock_drift_df.drop.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.isin.return_value = [True]
mock_drift_df.__getitem__.return_value = mock_drift_df
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = '2023-05-26 11:12:27+00:00'
mock_drift_df.__getitem__.return_value.dt.tz_localize.return_value.dt.strftime.return_value = (
'2023-05-26 11:12:27+00:00'
)
mock_drift_df.rename.return_value = mock_drift_df
mock_drift_df.drop_duplicates.return_value = mock_drift_df
mock_drift_df.to_dict.return_value = {'method': ['ks_test'], 'value': [0.5], 'feature': ['feature1'], 'timestamp': ['2023-05-26 11:12:27+00:00'], 'model_id': ['test_model_id'], 'accurate': [True]}
mock_drift_df.to_dict.return_value = [
{
'method': 'ks_test',
'value': 0.5,
'feature': 'feature1',
'timestamp': '2023-05-26 11:12:27+00:00',
'model_id': 'test_model_id',
'accurate': True,
}
]
model_metrics_activity.get_drift_metrics = AsyncMock(return_value=mock_drift_df)
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -458,16 +544,20 @@ async def test_calculate_drift_success_s(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert isinstance(result, dict)
assert result == mock_drift_df.to_dict.return_value
assert isinstance(result, list)
assert result == mock_drift_df.to_dict.return_value # type: ignore[comparison-overlap]
model_metrics_activity.info.assert_called()
model_metrics_activity.get_drift_metrics.assert_called_once()
# Verify transformations were called
mock_drift_df.drop.assert_called_once_with(columns=['p_value', 'chunk_start_date', 'chunk_end_date'], inplace=True)
mock_drift_df.drop.assert_called_once_with(columns=['p_value'], inplace=True)
mock_drift_df.__getitem__.assert_called()
mock_drift_df.rename.assert_called_once_with(columns={'metric': 'method', 'statistic': 'value'}, inplace=True)
mock_drift_df.drop_duplicates.assert_called_once_with(subset=['timestamp', 'method', 'feature'], keep='first', inplace=True)
mock_drift_df.to_dict.assert_called_once()
mock_drift_df.rename.assert_called_once_with(
columns={'metric': 'method', 'statistic': 'value'}, inplace=True
)
mock_drift_df.drop_duplicates.assert_called_once_with(
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
)
mock_drift_df.to_dict.assert_called_once_with(orient='records')
@mark.asyncio
@@ -478,15 +568,19 @@ async def test_calculate_drift_get_drift_metrics_error(
):
# Arrange
mock_to_datetime.return_value.dt.strftime.return_value = '2023-05-26 11:12:27'
model_metrics_activity.get_drift_metrics = AsyncMock(side_effect=Exception('Get drift metrics error'))
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
model_metrics_activity.get_drift_metrics = AsyncMock(
side_effect=Exception('Get drift metrics error')
)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
mock_target_df = MagicMock()
mock_target_df.pivot.return_value = mock_target_df
mock_target_df.index = ['2023-05-26 11:12:27']
@@ -515,10 +609,9 @@ async def test_calculate_drift_get_drift_metrics_error(
result = await model_metrics_activity.calculate_drift(input_data)
# Assert
assert result == {}
assert result == []
model_metrics_activity.error.assert_called_once_with(
'Error getting drift metrics: Get drift metrics error',
metadata['metadata']
'Error getting drift metrics: Get drift metrics error', metadata['metadata']
)
model_metrics_activity.send_notification_async.assert_called_once_with(
metadata=metadata['metadata'],
@@ -540,30 +633,36 @@ async def test_get_drift_metrics_success(
):
# Arrange
mock_time.return_value = 1000.0
mock_drift_df = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'metric': ['ks_test'],
'statistic': [0.5],
'feature': ['feature1'],
})
mock_drift_df = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'metric': ['ks_test'],
'statistic': [0.5],
'feature': ['feature1'],
}
)
mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock()
mock_model_analysis.return_value.get_drift_metrics_dataframe.return_value = mock_drift_df
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
@@ -596,21 +695,27 @@ async def test_get_drift_metrics_univariate_error(
):
# Arrange
mock_time.return_value = 1000.0
mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception('Univariate drift error')
reference_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
})
mock_model_analysis.return_value.detect_univariate_drift.side_effect = Exception(
'Univariate drift error'
)
reference_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27'],
'target': [1.0],
'feature1': [1.0],
}
)
reference_columns = reference_data.drop(
columns=['target', 'timestamp'], errors='ignore'
).columns
@@ -629,28 +734,26 @@ async def test_get_drift_metrics_univariate_error(
except Exception as e:
assert str(e) == 'Univariate drift error'
model_metrics_activity.error.assert_called_once_with(
'Error detecting univariate drift: Univariate drift error',
metadata['metadata']
'Error detecting univariate drift: Univariate drift error', metadata['metadata']
)
model_metrics_activity.emit_metric.assert_called_with(
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT,
tags=ANY
metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY
)
else:
raise AssertionError('Expected Exception')
@mark.asyncio
async def test_calculate_simple_metrics_success_all_metrics(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -673,23 +776,23 @@ async def test_calculate_simple_metrics_success_all_metrics(
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\', \'mse\', \'mae\', \'r2\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['rmse', 'mse', 'mae', 'r2']",
metadata['metadata'],
)
model_metrics_activity.debug.assert_called_once()
@mark.asyncio
async def test_calculate_simple_metrics_success_rmse_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_rmse_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -709,22 +812,21 @@ async def test_calculate_simple_metrics_success_rmse_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['rmse']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_mse_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_mse_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -744,22 +846,21 @@ async def test_calculate_simple_metrics_success_mse_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'mse\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['mse']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_mae_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_mae_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -779,22 +880,21 @@ async def test_calculate_simple_metrics_success_mae_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'mae\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['mae']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_r2_only(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_r2_only(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 2.0],
'prediction': [1.1, 2.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -814,23 +914,22 @@ async def test_calculate_simple_metrics_success_r2_only(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'r2\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_r2_zero_ss_tot(
model_metrics_activity
):
async def test_calculate_simple_metrics_r2_zero_ss_tot(model_metrics_activity):
# Arrange
# All target values are the same, so ss_tot will be 0
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 1.0],
'prediction': [1.1, 1.1],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'],
'target': [1.0, 1.0],
'prediction': [1.1, 1.1],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -851,22 +950,21 @@ async def test_calculate_simple_metrics_r2_zero_ss_tot(
assert result['data_size'].values[0] == 2
assert result['interval_minutes'].values[0] == 5
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'r2\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['r2']", metadata['metadata']
)
@mark.asyncio
async def test_calculate_simple_metrics_success_multiple_metrics_subset(
model_metrics_activity
):
async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_metrics_activity):
# Arrange
target_data = DataFrame({
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
})
target_data = DataFrame(
{
'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
'target': [1.0, 2.0, 3.0],
'prediction': [1.1, 2.1, 2.9],
}
)
input_data = {
**metadata,
'model_id': 'test_model_id',
@@ -887,7 +985,5 @@ async def test_calculate_simple_metrics_success_multiple_metrics_subset(
assert all(data_size == 3 for data_size in result['data_size'].values)
assert all(interval_minutes == 5 for interval_minutes in result['interval_minutes'].values)
model_metrics_activity.info.assert_called_once_with(
'Calculating simple metrics for model test_model_id: [\'rmse\', \'mae\']',
metadata['metadata']
"Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata']
)

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

15
tests/laborious/workflows/test_drift.py
View File

@@ -1,10 +1,10 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.drift import Drift
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
@@ -45,9 +45,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
reference_data = {'data': 'test_reference_data'}
drift_data = {'drift': 'test_drift_data'}
workflow_mock.start_local_activity_method.side_effect = [
target_data, reference_data
]
workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
workflow_mock.execute_local_activity_method.return_value = drift_data
workflow_mock.execute_activity_method = AsyncMock()
@@ -56,12 +54,13 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
await drift.run(input_data)
# Assert - Check start_local_activity_method calls
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
expected_gathering_query = f"""
SELECT *
FROM {input_data['schema']}.{input_data['source_table_name']}
FROM "{input_data['schema']}"."{input_data['source_table_name']}"
WHERE
model_id = {input_data['model_id']} AND
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
model_id = '{input_data['model_id']}' AND
timestamp > NOW() - INTERVAL {input_data['interval']} minutes
ORDER BY timestamp ASC
"""
@@ -73,6 +72,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift):
**metadata,
'query': expected_gathering_query,
'datetime_columns': ['timestamp', 'created_at'],
'orient': 'records',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -250,4 +250,3 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
retry_policy=ANY,
start_to_close_timeout=ANY,
)

25
tests/laborious/workflows/test_simple_metrics.py
View File

@@ -1,10 +1,10 @@
from unittest.mock import ANY, AsyncMock, call, patch
from pytest import fixture, mark
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from laborious.activities.activities import Activities
from laborious.workflows.simple_metrics import SimpleMetrics
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@fixture
@@ -42,9 +42,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
workflow_mock.execute_activity_method = AsyncMock()
@@ -52,17 +50,18 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
await simple_metrics.run(input_data)
# Assert - Check load_custom_query call
# Query format matches psycopg2.sql output (identifiers with double quotes, literals with single quotes)
expected_query = f"""
select p."timestamp", p.prediction, ld.value as "target"
from {input_data['schema']}.{input_data['predictions_table_name']} p
inner join {input_data['schema']}.{input_data['data_table_name']} ld
from "{input_data['schema']}"."{input_data['predictions_table_name']}" p
inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld
on p."timestamp" = ld."timestamp"
where
p.model_id = {input_data['model_id']} and
p.model_id = '{input_data['model_id']}' and
p.prediction is not null and
ld.variable = '{input_data['model_config']['target']}' and
ld.value is not null and
p."timestamp" >= NOW() - INTERVAL '{input_data['interval_minutes']} minutes'
p."timestamp" >= NOW() - INTERVAL {input_data['interval_minutes']} minutes
order by
p."timestamp" desc;
"""
@@ -75,6 +74,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics):
**metadata,
'query': expected_query,
'datetime_columns': ['timestamp'],
'orient': 'records',
},
retry_policy=ANY,
start_to_close_timeout=ANY,
@@ -159,9 +159,7 @@ async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics
target_data = {'data': 'test_target_data'}
simple_metrics_data = None
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
workflow_mock.execute_activity_method = AsyncMock()
@@ -193,9 +191,7 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim
target_data = {'data': 'test_target_data'}
simple_metrics_data = {'metrics': 'test_simple_metrics_data'}
workflow_mock.execute_local_activity_method.side_effect = [
target_data, simple_metrics_data
]
workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data]
workflow_mock.execute_activity_method = AsyncMock()
@@ -225,4 +221,3 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim
),
]
)