93d0849c80
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.
99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
from unittest.mock import AsyncMock, MagicMock, call, patch, ANY
|
|
from pytest import fixture, mark
|
|
from laborious.activities.activities import Activities
|
|
from laborious.workflows.minimal_retrain import MinimalRetrain
|
|
|
|
|
|
@fixture
|
|
def minimal_retrain() -> MinimalRetrain:
|
|
return MinimalRetrain()
|
|
|
|
|
|
metadata = {
|
|
"metadata": {
|
|
"model_id": "test_model_id",
|
|
"model_name": "test_model",
|
|
"workflow_name": "minimal_retrain",
|
|
"schedule_name": "test_schedule",
|
|
},
|
|
}
|
|
|
|
|
|
@mark.asyncio
|
|
@patch('laborious.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
|
|
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
|
|
input_data = {
|
|
"model_id": "test_model_id",
|
|
"model_name": "test_model",
|
|
"workflow_name": "minimal_retrain",
|
|
"schedule_name": "test_schedule",
|
|
"query": "test_query",
|
|
"schema": "test_schema",
|
|
"table_name": "test_table",
|
|
}
|
|
|
|
workflow_mock.execute_activity_method = AsyncMock(
|
|
return_value={
|
|
"data1": "1",
|
|
"data2": "2",
|
|
}
|
|
)
|
|
|
|
await minimal_retrain.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
|
|
)
|
|
]
|
|
)
|
|
|
|
workflow_mock.execute_activity_method.assert_has_calls([
|
|
call(
|
|
Activities.retrain_model,
|
|
{
|
|
**metadata,
|
|
'data': workflow_mock.execute_local_activity_method.return_value,
|
|
'model_name': input_data['model_name'],
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY
|
|
)
|
|
])
|
|
|
|
workflow_mock.execute_activity_method.assert_has_calls([
|
|
call(
|
|
Activities.update_production_model,
|
|
{
|
|
**metadata,
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id'],
|
|
**workflow_mock.execute_activity_method.return_value,
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY
|
|
)
|
|
])
|
|
|
|
workflow_mock.execute_activity_method.assert_has_calls([
|
|
call(
|
|
Activities.export_data_to_postgres,
|
|
{
|
|
**metadata,
|
|
'data': workflow_mock.execute_activity_method.return_value,
|
|
'schema': input_data['schema'],
|
|
'table_name': input_data['table_name'],
|
|
},
|
|
retry_policy=ANY,
|
|
start_to_close_timeout=ANY
|
|
)
|
|
])
|