From d22545ba881b7197cb6133828b30a2f8edbabedb Mon Sep 17 00:00:00 2001
From: vitor-aignosi <vitor.santos@aignosi.com.br>
Date: Mon, 12 May 2025 14:50:20 -0300
Subject: [PATCH] SIENTIAPDE-994

Finishing base pipelines, preparing to testing
---
 laborious/activities/postgres.py              |  2 +-
 .../subworkflows/test_predictions_batch.py    | 48 +++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 tests/laborious/workflows/subworkflows/test_predictions_batch.py

diff --git a/laborious/activities/postgres.py b/laborious/activities/postgres.py
index 3716fa9..7d8c7f6 100644
--- a/laborious/activities/postgres.py
+++ b/laborious/activities/postgres.py
@@ -40,7 +40,7 @@ class Postgres(BaseActivity):
         self.close()
 
     @activity.defn(name="load_custom_query")
-    async def load_custom_query(self, query: str) -> dict[str, dict]:
+    async def load_custom_query(self, query: str) -> dict[str, Any]:
         """
         Loads data from a custom query.
 
diff --git a/tests/laborious/workflows/subworkflows/test_predictions_batch.py b/tests/laborious/workflows/subworkflows/test_predictions_batch.py
new file mode 100644
index 0000000..065ca31
--- /dev/null
+++ b/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)
+    ])