vitor-aignosi
64f0747e8e
Update version and enhance metrics calculation in Laborious system - Updated image tag in values.yaml from 1.1.0 to 1.1.1. - Modified GITHUB_BRANCH environment variable for consistency. - Added a new method `calculate_simple_metrics` in model_metrics.py to compute various model performance metrics including RMSE, MSE, MAE, and R2. - Integrated the new metrics calculation into the worker setup, allowing for concurrent processing of simple metrics. - Updated tests to cover the new metrics calculation functionality, ensuring comprehensive validation of the implementation.
254 lines
8.1 KiB
Python
254 lines
8.1 KiB
Python
from unittest.mock import ANY, AsyncMock, call, patch
|
|
|
|
from pytest import fixture, mark
|
|
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.drift import Drift
|
|
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
|
|
|
|
|
@fixture
|
|
def drift() -> Drift:
|
|
return Drift()
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model_id',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'drift',
|
|
'schedule_name': 'test_schedule',
|
|
},
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
|
async def test_run(workflow_mock: AsyncMock, drift: Drift):
|
|
# Arrange
|
|
input_data = {
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
'schema': 'test_schema',
|
|
'source_table_name': 'test_source_table',
|
|
'target_table_name': 'test_target_table',
|
|
'interval': 60,
|
|
'model_config': {'target': 'test_target'},
|
|
'drift_metrics': ['psi', 'ks'],
|
|
'chunk_period': 'hour',
|
|
}
|
|
|
|
target_name = input_data['model_config']['target']
|
|
|
|
target_data = {'data': 'test_target_data'}
|
|
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.execute_local_activity_method.return_value = drift_data
|
|
workflow_mock.execute_activity_method = AsyncMock()
|
|
|
|
# Act
|
|
await drift.run(input_data)
|
|
|
|
# Assert - Check start_local_activity_method calls
|
|
expected_gathering_query = f"""
|
|
SELECT *
|
|
FROM {input_data['schema']}.{input_data['source_table_name']}
|
|
WHERE
|
|
model_id = {input_data['model_id']} AND
|
|
timestamp > NOW() - INTERVAL '{input_data['interval']} minutes'
|
|
ORDER BY timestamp ASC
|
|
"""
|
|
|
|
workflow_mock.start_local_activity_method.assert_has_calls(
|
|
[
|
|
call(
|
|
Activities.load_custom_query,
|
|
{
|
|
**metadata,
|
|
'query': expected_gathering_query,
|
|
'datetime_columns': ['timestamp', 'created_at'],
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
),
|
|
call(
|
|
Activities.get_reference_data,
|
|
{
|
|
**metadata,
|
|
'model_name': input_data['model_name'],
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
),
|
|
]
|
|
)
|
|
|
|
# Assert - Check calculate_drift call
|
|
workflow_mock.execute_local_activity_method.assert_called_once_with(
|
|
Activities.calculate_drift,
|
|
{
|
|
**metadata,
|
|
'target_data': target_data,
|
|
'reference_data': reference_data,
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id'],
|
|
'target_name': target_name,
|
|
'drift_metrics': input_data['drift_metrics'],
|
|
'chunk_period': input_data['chunk_period'],
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
|
|
# Assert - Check export_data_to_postgres call
|
|
workflow_mock.execute_activity_method.assert_called_once_with(
|
|
Activities.export_data_to_postgres,
|
|
{
|
|
**metadata,
|
|
'data': drift_data,
|
|
'schema': input_data['schema'],
|
|
'table_name': input_data['target_table_name'],
|
|
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
|
async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift):
|
|
# Arrange
|
|
input_data = {
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
'schema': 'test_schema',
|
|
'source_table_name': 'test_source_table',
|
|
'target_table_name': 'test_target_table',
|
|
'interval': 60,
|
|
'model_config': {'target': 'test_target'},
|
|
'drift_metrics': ['psi', 'ks'],
|
|
}
|
|
|
|
target_data = None
|
|
reference_data = {'data': 'test_reference_data'}
|
|
|
|
workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data]
|
|
|
|
workflow_mock.execute_local_activity_method = AsyncMock()
|
|
workflow_mock.execute_activity_method = AsyncMock()
|
|
|
|
# Act
|
|
await drift.run(input_data)
|
|
|
|
# Assert - Should not call calculate_drift or export
|
|
workflow_mock.execute_local_activity_method.assert_not_called()
|
|
workflow_mock.execute_activity_method.assert_not_called()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
|
async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift):
|
|
# Arrange
|
|
input_data = {
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
'schema': 'test_schema',
|
|
'source_table_name': 'test_source_table',
|
|
'target_table_name': 'test_target_table',
|
|
'interval': 60,
|
|
'model_config': {'target': 'test_target'},
|
|
'drift_metrics': ['psi', 'ks'],
|
|
}
|
|
|
|
target_name = input_data['model_config']['target']
|
|
|
|
target_data = {'data': 'test_target_data'}
|
|
reference_data = {'data': 'test_reference_data'}
|
|
drift_data = None
|
|
|
|
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()
|
|
|
|
# Act
|
|
await drift.run(input_data)
|
|
|
|
# Assert - Should call calculate_drift but not export
|
|
workflow_mock.execute_local_activity_method.assert_called_once_with(
|
|
Activities.calculate_drift,
|
|
{
|
|
**metadata,
|
|
'target_data': target_data,
|
|
'reference_data': reference_data,
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id'],
|
|
'target_name': target_name,
|
|
'drift_metrics': input_data['drift_metrics'],
|
|
'chunk_period': input_data.get('chunk_period', 'min'),
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
|
|
workflow_mock.execute_activity_method.assert_not_called()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.workflows.drift.workflow', new_callable=AsyncMock)
|
|
async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift):
|
|
# Arrange
|
|
input_data = {
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
'schema': 'test_schema',
|
|
'source_table_name': 'test_source_table',
|
|
'target_table_name': 'test_target_table',
|
|
'interval': 60,
|
|
'model_config': {'target': 'test_target'},
|
|
'drift_metrics': ['psi', 'ks'],
|
|
# chunk_period not provided, should default to 'min'
|
|
}
|
|
|
|
target_name = input_data['model_config']['target']
|
|
|
|
target_data = {'data': 'test_target_data'}
|
|
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.execute_local_activity_method.return_value = drift_data
|
|
workflow_mock.execute_activity_method = AsyncMock()
|
|
|
|
# Act
|
|
await drift.run(input_data)
|
|
|
|
# Assert - Check calculate_drift call with default chunk_period
|
|
workflow_mock.execute_local_activity_method.assert_called_once_with(
|
|
Activities.calculate_drift,
|
|
{
|
|
**metadata,
|
|
'target_data': target_data,
|
|
'reference_data': reference_data,
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id'],
|
|
'target_name': target_name,
|
|
'drift_metrics': input_data['drift_metrics'],
|
|
'chunk_period': 'min', # Default value
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
|