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, )