78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
from unittest.mock import ANY, AsyncMock, call, patch
|
|
|
|
from pytest import fixture, mark
|
|
|
|
from model_manager.activities.activities import Activities
|
|
from model_manager.workflows.predictions_batch import PredictionsBatch
|
|
|
|
|
|
@fixture
|
|
def predictions_batch() -> PredictionsBatch:
|
|
return PredictionsBatch()
|
|
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'model_id': 'test_model_id',
|
|
'model_name': 'test_model',
|
|
'workflow_name': 'predictions_batch',
|
|
'schedule_name': 'test_schedule',
|
|
},
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('model_manager.workflows.predictions_batch.workflow', new_callable=AsyncMock)
|
|
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
|
|
workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'}
|
|
input_data = {
|
|
'schedule_name': 'test_schedule',
|
|
'model_name': 'test_model',
|
|
'model_id': 'test_model_id',
|
|
'query': 'SELECT * FROM test',
|
|
'schema': 'test_schema',
|
|
'table_name': 'test_table',
|
|
'datetime_columns': ['timestamp', 'created_at'],
|
|
'prediction_store_policy': 'erl:1',
|
|
'model_config': {'retention': '30'},
|
|
}
|
|
|
|
await predictions_batch.run(input_data)
|
|
|
|
workflow_mock.execute_local_activity_method.assert_has_calls(
|
|
[
|
|
call(
|
|
Activities.load_custom_query,
|
|
{
|
|
**metadata,
|
|
'query': input_data['query'],
|
|
'datetime_columns': input_data.get('datetime_columns', []),
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY,
|
|
)
|
|
]
|
|
)
|
|
prediction_input = {
|
|
'metadata': metadata,
|
|
'data': {'data': 'test_data'},
|
|
'schema': input_data['schema'],
|
|
'table_name': input_data['table_name'],
|
|
'model_id': input_data['model_id'],
|
|
'model_name': input_data['model_name'],
|
|
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
|
|
'mlflow_transform_filters': input_data.get(
|
|
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
|
),
|
|
'mlflow_predict_filters': input_data.get(
|
|
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
|
),
|
|
'model_config': input_data.get('model_config', {}),
|
|
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
|
|
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'),
|
|
}
|
|
|
|
workflow_mock.execute_child_workflow.assert_has_calls(
|
|
[call('prediction_process', prediction_input)]
|
|
)
|