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.
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 model_manager.activities.activities import Activities
|
|
from model_manager.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('model_manager.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
|
|
)
|
|
])
|