vitor-aignosi
503d9aa485
Update dependencies and refactor input filter handling for consistency - Updated sientia-dataops-library dependency version from 1.10.3 to 1.10.4 in requirements.txt. - Refactored input filter handling in the Gates class to read policy and config keys in a case-insensitive manner. - Updated test cases to ensure consistency in filter key naming conventions across various scenarios.
249 lines
8.2 KiB
Python
249 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_activity_method.side_effect = [target_data, reference_data]
|
|
workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
|
|
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
|
|
|
|
# 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_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_activity_method.side_effect = [target_data, reference_data]
|
|
|
|
workflow_mock.execute_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_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_activity_method.side_effect = [target_data, reference_data]
|
|
workflow_mock.execute_local_activity_method = AsyncMock(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_activity_method.side_effect = [target_data, reference_data]
|
|
workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
|
|
workflow_mock.execute_activity_method = AsyncMock(return_value=None)
|
|
|
|
# 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,
|
|
)
|