Code import - branch feature/SIENTIAPDE-1646
This commit is contained in:
397
tests/workflow/sub_workflows/test_core_scouter.py
Normal file
397
tests/workflow/sub_workflows/test_core_scouter.py
Normal file
@@ -0,0 +1,397 @@
|
||||
from unittest.mock import ANY, AsyncMock, call, patch
|
||||
|
||||
import pytest
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
|
||||
from scouter.activities.activities import Activities
|
||||
from scouter.workflow.sub_workflows.core_scouter import CoreScouter
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def core_scouter():
|
||||
return CoreScouter()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('scouter.workflow.sub_workflows.core_scouter.workflow', new_callable=AsyncMock)
|
||||
async def test_core_scouter_workflow_success(mock_workflow, core_scouter):
|
||||
mock_workflow.execute_local_activity_method.side_effect = [
|
||||
'filtered_data',
|
||||
'grouped_data',
|
||||
'held_data',
|
||||
]
|
||||
mock_workflow.execute_activity_method.side_effect = [
|
||||
{'affected_rows': 10}, # export_data_to_postgres
|
||||
None, # write_metrics
|
||||
None, # store_data_package
|
||||
]
|
||||
await core_scouter.run(
|
||||
input_data={
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
},
|
||||
'workflow_name': 'test_workflow',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'data': 'test_data',
|
||||
'trigger_laborious': False,
|
||||
'filters': {'test_filter': 'test_value'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'retention_time': 3600,
|
||||
'model_tags': {},
|
||||
'debug_data_package': True,
|
||||
'fill_missing_tags': False,
|
||||
}
|
||||
)
|
||||
|
||||
expected_metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
}
|
||||
|
||||
mock_workflow.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.data_quality_gate,
|
||||
{
|
||||
**expected_metadata,
|
||||
'filters': {'test_filter': 'test_value'},
|
||||
'data': 'test_data',
|
||||
'model_tags': {},
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
mock_workflow.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.aggregate_data,
|
||||
{**expected_metadata, 'data': 'filtered_data', 'model_tags': {}},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
mock_workflow.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.group_and_hold_data,
|
||||
{
|
||||
**expected_metadata,
|
||||
'workflow_name': 'test_workflow',
|
||||
'schedule_name': 'test_schedule',
|
||||
'data': 'grouped_data',
|
||||
'model_id': 'test_model_id',
|
||||
'retention_time': 3600,
|
||||
'model_tags': {},
|
||||
'fill_missing_tags': False,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
mock_workflow.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**expected_metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'data': 'held_data',
|
||||
'timestamp_conversion': {
|
||||
'column': 'timestamp',
|
||||
'format': DATETIME_FORMAT_WITH_TZ,
|
||||
},
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
mock_workflow.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.write_metrics,
|
||||
{
|
||||
**expected_metadata,
|
||||
'tag_values': 'held_data',
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
mock_workflow.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.store_data_package,
|
||||
{
|
||||
**expected_metadata,
|
||||
'workflow_name': 'test_workflow',
|
||||
'schedule_name': 'test_schedule',
|
||||
'held_data': 'held_data',
|
||||
'data': 'test_data',
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('scouter.workflow.sub_workflows.core_scouter.workflow', new_callable=AsyncMock)
|
||||
async def test_core_scouter_workflow_with_empty_data(mock_workflow, core_scouter):
|
||||
mock_workflow.execute_local_activity_method.return_value = {}
|
||||
await core_scouter.run(
|
||||
input_data={
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
},
|
||||
'workflow_name': 'test_workflow',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'data': 'test_data',
|
||||
'trigger_laborious': False,
|
||||
'filters': {'test_filter': 'test_value'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'retention_time': 3600,
|
||||
'model_tags': {},
|
||||
'fill_missing_tags': False,
|
||||
}
|
||||
)
|
||||
|
||||
expected_metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
}
|
||||
|
||||
mock_workflow.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.data_quality_gate,
|
||||
{
|
||||
**expected_metadata,
|
||||
'filters': {'test_filter': 'test_value'},
|
||||
'data': 'test_data',
|
||||
'model_tags': {},
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
mock_workflow.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.aggregate_data,
|
||||
{**expected_metadata, 'data': {}, 'model_tags': {}},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
mock_workflow.execute_local_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.group_and_hold_data,
|
||||
{
|
||||
**expected_metadata,
|
||||
'workflow_name': 'test_workflow',
|
||||
'schedule_name': 'test_schedule',
|
||||
'data': {},
|
||||
'model_id': 'test_model_id',
|
||||
'retention_time': 3600,
|
||||
'model_tags': {},
|
||||
'fill_missing_tags': False,
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert mock_workflow.execute_local_activity_method.call_count == 3
|
||||
mock_workflow.execute_activity_method.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('scouter.workflow.sub_workflows.core_scouter.workflow', new_callable=AsyncMock)
|
||||
async def test_core_scouter_workflow_with_zero_affected_rows(mock_workflow, core_scouter):
|
||||
"""
|
||||
Test that workflow stops after export when no rows are affected
|
||||
"""
|
||||
mock_workflow.execute_local_activity_method.side_effect = [
|
||||
'filtered_data',
|
||||
'grouped_data',
|
||||
'held_data',
|
||||
]
|
||||
mock_workflow.execute_activity_method.return_value = {'affected_rows': 0}
|
||||
|
||||
await core_scouter.run(
|
||||
input_data={
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
},
|
||||
'workflow_name': 'test_workflow',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'data': 'test_data',
|
||||
'trigger_laborious': False,
|
||||
'filters': {'test_filter': 'test_value'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'retention_time': 3600,
|
||||
'model_tags': {},
|
||||
'debug_data_package': True,
|
||||
'fill_missing_tags': False,
|
||||
}
|
||||
)
|
||||
|
||||
expected_metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
}
|
||||
|
||||
mock_workflow.execute_activity_method.assert_called_once_with(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**expected_metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'data': 'held_data',
|
||||
'timestamp_conversion': {
|
||||
'column': 'timestamp',
|
||||
'format': DATETIME_FORMAT_WITH_TZ,
|
||||
},
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('scouter.workflow.sub_workflows.core_scouter.workflow', new_callable=AsyncMock)
|
||||
async def test_core_scouter_workflow_without_debug_data_package(mock_workflow, core_scouter):
|
||||
"""
|
||||
Test that store_data_package is not called when debug_data_package is False
|
||||
"""
|
||||
mock_workflow.execute_local_activity_method.side_effect = [
|
||||
'filtered_data',
|
||||
'grouped_data',
|
||||
'held_data',
|
||||
]
|
||||
mock_workflow.execute_activity_method.side_effect = [
|
||||
{'affected_rows': 5},
|
||||
None,
|
||||
]
|
||||
|
||||
await core_scouter.run(
|
||||
input_data={
|
||||
'metadata': {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
},
|
||||
'workflow_name': 'test_workflow',
|
||||
'schedule_name': 'test_schedule',
|
||||
'model_name': 'test_model',
|
||||
'model_id': 'test_model_id',
|
||||
'data': 'test_data',
|
||||
'trigger_laborious': False,
|
||||
'filters': {'test_filter': 'test_value'},
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'retention_time': 3600,
|
||||
'model_tags': {},
|
||||
'debug_data_package': False,
|
||||
'fill_missing_tags': False,
|
||||
}
|
||||
)
|
||||
|
||||
expected_metadata = {
|
||||
'metadata': {
|
||||
'model_id': 'test_model_id',
|
||||
'model_name': 'test_model',
|
||||
'schedule_name': 'test_schedule',
|
||||
'workflow_name': 'test_workflow',
|
||||
}
|
||||
}
|
||||
|
||||
mock_workflow.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.export_data_to_postgres,
|
||||
{
|
||||
**expected_metadata,
|
||||
'schema': 'test_schema',
|
||||
'table_name': 'test_table',
|
||||
'data': 'held_data',
|
||||
'timestamp_conversion': {
|
||||
'column': 'timestamp',
|
||||
'format': DATETIME_FORMAT_WITH_TZ,
|
||||
},
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
mock_workflow.execute_activity_method.assert_has_calls(
|
||||
[
|
||||
call(
|
||||
Activities.write_metrics,
|
||||
{
|
||||
**expected_metadata,
|
||||
'tag_values': 'held_data',
|
||||
},
|
||||
retry_policy=ANY,
|
||||
start_to_close_timeout=ANY,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert mock_workflow.execute_activity_method.call_count == 2
|
||||
Reference in New Issue
Block a user