SIENTIAPDE-994

Finishing base pipelines, preparing to testing

tests/laborious/workflows/subworkflows/test_predictions_batch.py

@@ -0,0 +1,48 @@
+from unittest.mock import AsyncMock, call, patch
+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()
+
+
+@mark.asyncio
+@patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock)
+async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
+    workflow_mock.execute_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'
+    }
+
+    await predictions_batch.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_has_calls([
+        call(
+            Activities.prepare_activity,
+            {
+                'schedule_name': input_data['schedule_name'],
+                'model_name': input_data['model_name'],
+                'model_id': input_data['model_id']
+            }
+        )
+    ])
+
+    workflow_mock.execute_activity_method.assert_has_calls([
+        call(
+            Activities.load_custom_query,
+            input_data['query']
+        )
+    ])
+
+    workflow_mock.execute_child_workflow.assert_has_calls([
+        call(
+            'prediction_process', input_data)
+    ])