vitor-aignosi
7e4e048ca3
Update sientia-mlops-library dependency to version 0.40.5 and fix SQL interval formatting in Drift and SimpleMetrics workflows for improved query accuracy.
253 lines
8.2 KiB
Python
253 lines
8.2 KiB
Python
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
|
|
|
|
|
|
@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
|
|
# 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']}"
|
|
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'],
|
|
'orient': 'records',
|
|
},
|
|
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,
|
|
)
|