diff --git a/tests/orchestrator/activities/test_formatters.py b/tests/orchestrator/activities/test_formatters.py index b75cdfc..04ada97 100644 --- a/tests/orchestrator/activities/test_formatters.py +++ b/tests/orchestrator/activities/test_formatters.py @@ -34,7 +34,9 @@ metadata = { return_value={"test_scouter": "test_scouter"}) @patch("orchestrator.activities.formatters.predictions_batch", return_value={"test_predictions_batch": "test_predictions_batch"}) -async def test_process_schedules(mock_predictions_batch, mock_scouter, formatters): +@patch("orchestrator.activities.formatters.minimal_retrain", + return_value={"test_minimal_retrain": "test_minimal_retrain"}) +async def test_process_schedules(mock_minimal_retrain, mock_predictions_batch, mock_scouter, formatters): input_data = { "pipelines": [ { @@ -50,6 +52,13 @@ async def test_process_schedules(mock_predictions_batch, mock_scouter, formatter "model_name": "test_model_name", "model_id": "test_model_id", "updated_at": "2021-01-02" + }, + { + "schedule_name": "test_schedule_name3", + "workflow_type": "minimal_retrain", + "model_name": "test_model_name", + "model_id": "test_model_id", + "updated_at": "2021-01-03" } ] } @@ -67,6 +76,10 @@ async def test_process_schedules(mock_predictions_batch, mock_scouter, formatter "test_schedule_name2": { "test_predictions_batch": "test_predictions_batch", "updated_at": "2021-01-02" + }, + "test_schedule_name3": { + "test_minimal_retrain": "test_minimal_retrain", + "updated_at": "2021-01-03" } } } diff --git a/tests/orchestrator/utils/test_orchestrator_functions.py b/tests/orchestrator/utils/test_orchestrator_functions.py index e0b08a9..eb4b056 100644 --- a/tests/orchestrator/utils/test_orchestrator_functions.py +++ b/tests/orchestrator/utils/test_orchestrator_functions.py @@ -1,6 +1,7 @@ from unittest.mock import patch, call from orchestrator.utils.orchestrator_functions import ( common_config, + minimal_retrain, scouter, predictions_batch, overlap_filter_config, @@ -31,6 +32,31 @@ def test_common_config(): assert result == expected +def test_minimal_retrain(): + config = { + "workflow_type": "minimal_retrain", + "schedule_name": "test_schedule", + "model_id": "test_model_id", + "models": { + "name": "test_model_name" + }, + "query": "select * from sientia_data.laborious_data order by \"timestamp\" desc limit 30;", + } + result = minimal_retrain(config) + expected = { + "workflow_type": "minimal_retrain", + "schedule_name": "test_schedule", + "frequency": "1m", + "max_retry_policy": 1, + "model_id": "test_model_id", + "model_name": "test_model_name", + "query": "select * from sientia_data.laborious_data order by \"timestamp\" desc limit 30;", + "schema": "sientia_data", + "table_name": "log_retrain", + } + assert result == expected + + def test_scouter(): config = { "workflow_type": "scouter",