SIENTIAPDE-1243: Initial commit of the model manager project, adding core files and configurations.
This commit introduces the initial project structure, including: - .env.example: Example environment configuration. - .github/workflows/quality-gate.yml: CI workflow for quality checks. - .gitignore: Specifies intentionally untracked files that Git should ignore. - Makefile: Automation of tasks like docker builds. - README.md: Project documentation. - Source code for model management, activities, utils, worker and workflows. - Test suite. - Dockerfile for the simulator. - sonar-project.properties: SonarQube configuration file. - values.yaml: Helm chart values for deployment.
This commit is contained in:
88
tests/laborious/workflows/test_predictions_batch.py
Normal file
88
tests/laborious/workflows/test_predictions_batch.py
Normal file
@@ -0,0 +1,88 @@
|
||||
from unittest.mock import AsyncMock, call, patch, ANY
|
||||
from pytest import fixture, mark
|
||||
from laborious.activities.activities import Activities
|
||||
from laborious.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('laborious.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)
|
||||
])
|
||||
Reference in New Issue
Block a user