aba4a3f1a5
This commit renames the 'laborious' package to 'model_manager' across the entire project. This includes renaming directories, modules, references in code, configuration files, and documentation to reflect the new package name. This change improves clarity and consistency within the project.
89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
from unittest.mock import AsyncMock, call, patch, ANY
|
|
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',
|
|
'opc_output_config': 'test_opc_output_config',
|
|
'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']),
|
|
'opc_output_config': input_data.get('opc_output_config', {}),
|
|
'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1')
|
|
}
|
|
|
|
workflow_mock.execute_child_workflow.assert_has_calls([
|
|
call(
|
|
'prediction_process', prediction_input)
|
|
])
|