vitor-aignosi
b68674fe64
Enhance MLFlowRepository and Activities classes with new methods and metrics - Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry. - Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models. - Updated Activities class to integrate ModelMetrics for improved metrics handling. - Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities. - Updated various workflows to include `transform_table_name` in input data for better data handling.
248 lines
7.9 KiB
Python
248 lines
7.9 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,
|
|
'target_name': 'test_target',
|
|
'drift_metrics': ['psi', 'ks'],
|
|
'chunk_period': 'hour',
|
|
}
|
|
|
|
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': input_data['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,
|
|
'target_name': '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,
|
|
'target_name': 'test_target',
|
|
'drift_metrics': ['psi', 'ks'],
|
|
}
|
|
|
|
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': input_data['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,
|
|
'target_name': 'test_target',
|
|
'drift_metrics': ['psi', 'ks'],
|
|
# chunk_period not provided, should default to 'min'
|
|
}
|
|
|
|
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': input_data['target_name'],
|
|
'drift_metrics': input_data['drift_metrics'],
|
|
'chunk_period': 'min', # Default value
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
|