130 lines
3.7 KiB
Python
130 lines
3.7 KiB
Python
from unittest.mock import ANY, AsyncMock, call, patch
|
|
|
|
from pytest import fixture, mark
|
|
|
|
from scouter.activities.activities import Activities
|
|
from scouter.workflow.scouter import Scouter
|
|
|
|
|
|
@fixture
|
|
def scouter():
|
|
return Scouter()
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('scouter.workflow.scouter.workflow', new_callable=AsyncMock)
|
|
async def test_scouter_workflow(mock_workflow, scouter):
|
|
mock_workflow.execute_local_activity_method.side_effect = [
|
|
'test_last_data_timestamp',
|
|
'test_data',
|
|
]
|
|
await scouter.run(
|
|
input_data={
|
|
'topic': 'test_topic',
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
}
|
|
)
|
|
|
|
expected_metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model_id',
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test_schedule',
|
|
'workflow_name': 'scouter',
|
|
}
|
|
}
|
|
|
|
mock_workflow.execute_local_activity_method.assert_has_calls(
|
|
[
|
|
call(
|
|
Activities.get_last_data_timestamp,
|
|
{**expected_metadata, 'workflow_name': 'scouter', 'schedule_name': 'test_schedule'},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
]
|
|
)
|
|
|
|
mock_workflow.execute_local_activity_method.assert_has_calls(
|
|
[
|
|
call(
|
|
Activities.load_latest_data,
|
|
{
|
|
**expected_metadata,
|
|
'collection_name': 'raw_test_schedule',
|
|
'last_data_timestamp': 'test_last_data_timestamp',
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
]
|
|
)
|
|
|
|
mock_workflow.execute_activity_method.assert_called_once_with(
|
|
Activities.put_last_data_timestamp,
|
|
{
|
|
**expected_metadata,
|
|
'data': 'test_data',
|
|
'workflow_name': 'scouter',
|
|
'schedule_name': 'test_schedule',
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
|
|
mock_workflow.execute_child_workflow.assert_called_once_with(
|
|
'subworkflow.core_scouter',
|
|
{
|
|
'metadata': expected_metadata,
|
|
'topic': 'test_topic',
|
|
'data': 'test_data',
|
|
'workflow_name': 'scouter',
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
},
|
|
)
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('scouter.workflow.scouter.workflow', new_callable=AsyncMock)
|
|
async def test_scouter_workflow_empty(mock_workflow, scouter):
|
|
mock_workflow.execute_local_activity_method.side_effect = ['test_last_data_timestamp', {}]
|
|
await scouter.run(
|
|
input_data={
|
|
'topic': 'test_topic',
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
}
|
|
)
|
|
|
|
expected_metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model_id',
|
|
'model_name': 'test_model',
|
|
'schedule_name': 'test_schedule',
|
|
'workflow_name': 'scouter',
|
|
}
|
|
}
|
|
|
|
mock_workflow.execute_local_activity_method.assert_has_calls(
|
|
[
|
|
call(
|
|
Activities.load_latest_data,
|
|
{
|
|
**expected_metadata,
|
|
'collection_name': 'raw_test_schedule',
|
|
'last_data_timestamp': 'test_last_data_timestamp',
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
]
|
|
)
|
|
|
|
mock_workflow.execute_activity_method.assert_not_called()
|
|
mock_workflow.execute_child_workflow.assert_not_called()
|