SIENTIAPDE-1712

Refactor MinioDataFramePayload usage across activities

- Updated instances of MinioDataFramePayload initialization in Gates, MLFlow, and Storage classes to use the new from_dict method for better data reconstruction from dictionaries.
- Enhanced the PredictionProcess workflow to utilize the updated payload handling.
- Added passthrough fixtures in tests to accommodate the new from_dict method for consistent testing behavior.

tests/laborious/activities/test_gates.py

@@ -7,6 +7,14 @@ from sientia_do.notifications.models import NotificationLevel
 from laborious.activities.gates import Gates
 
 
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.activities.gates.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+    ):
+        yield
+
+
 def _minio_payload(retrieve_return, status=None):
     """
     Build a MinioDataFramePayload-like test double with async retrieve.

tests/laborious/activities/test_mlflow.py

@@ -8,6 +8,14 @@ from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_
 from laborious.activities.mlflow import MLFlow
 
 
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+    ):
+        yield
+
+
 @patch('laborious.activities.mlflow.MLFlowRepository')
 @patch('laborious.activities.mlflow.MinioRepository')
 def test___init__(mock_minio_repository, mock_mlflow_repository):

tests/laborious/activities/test_storage.py

@@ -8,6 +8,15 @@ from sientia_do.temporal.activities.postgres import Postgres
 
 from laborious.activities.storage import Storage
 
+
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.activities.storage.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+    ):
+        yield
+
+
 metadata = {
     'metadata': {
         'model_id': 'test_model_id',

tests/laborious/utils/models/test_minio_dataframe_payload.py

@@ -214,3 +214,55 @@ async def test_from_dataframe_offloaded(mock_now):
     assert result.bucket == 'test-bucket'
     assert result.uri == 's3://test-bucket/full/key.parquet'
     minio.upload_file.assert_awaited_once()
+
+
+def test_from_dict_inline():
+    raw = {
+        'last_timestamp': '2024-01-01T00:00:00+00:00',
+        'status': None,
+        'data': {'col1': {0: 'val1'}},
+        'bucket': None,
+        'object_key': None,
+        'object_prefix': None,
+        'uri': None,
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.last_timestamp == '2024-01-01T00:00:00+00:00'
+    assert payload.data == {'col1': {0: 'val1'}}
+    assert payload.object_key is None
+
+
+def test_from_dict_offloaded():
+    raw = {
+        'last_timestamp': '2024-06-15T10:30:45+00:00',
+        'status': {'success': True},
+        'data': None,
+        'bucket': 'my-bucket',
+        'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+        'object_prefix': 'training_datasets/model',
+        'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet',
+    }
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert isinstance(payload, MinioDataFramePayload)
+    assert payload.data is None
+    assert payload.bucket == 'my-bucket'
+    assert payload.object_key == raw['object_key']
+    assert payload.object_prefix == 'training_datasets/model'
+    assert payload.uri == raw['uri']
+    assert payload.status == {'success': True}
+
+
+def test_from_dict_minimal_keys():
+    raw = {'last_timestamp': '2024-01-01'}
+    payload = MinioDataFramePayload.from_dict(raw)
+    assert payload.last_timestamp == '2024-01-01'
+    assert payload.data is None
+    assert payload.bucket is None
+    assert payload.object_key is None
+
+
+def test_from_dict_passthrough_existing_instance():
+    original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b')
+    result = MinioDataFramePayload.from_dict(original)
+    assert result is original

tests/laborious/workflows/subworkflows/test_prediction_process.py

@@ -6,6 +6,15 @@ from laborious.activities.activities import Activities
 from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
 
 
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.workflows.sub_workflows.prediction_process.MinioDataFramePayload.from_dict',
+        side_effect=lambda x: x,
+    ):
+        yield
+
+
 @fixture
 def prediction_process():
     return PredictionProcess()

tests/laborious/workflows/test_minimal_retrain.py

@@ -6,6 +6,15 @@ from laborious.activities.activities import Activities
 from laborious.workflows.minimal_retrain import MinimalRetrain
 
 
+@fixture(autouse=True)
+def _passthrough_from_dict():
+    with patch(
+        'laborious.workflows.minimal_retrain.MinioDataFramePayload.from_dict',
+        side_effect=lambda x: x,
+    ):
+        yield
+
+
 @fixture
 def minimal_retrain() -> MinimalRetrain:
     return MinimalRetrain()

tests/laborious/workflows/test_predictions_batch.py

@@ -1,4 +1,4 @@
-from unittest.mock import ANY, AsyncMock, call, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
 
 from pytest import fixture, mark
 
@@ -22,12 +22,15 @@ metadata = {
 
 
 @mark.asyncio
+@patch(
+    'laborious.workflows.predictions_batch.MinioDataFramePayload.from_dict', side_effect=lambda x: x
+)
 @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 = {
-        'success': True,
-        'data': {'col': ['test_data']},
-    }
+async def test_run(workflow_mock: AsyncMock, mock_from_dict, predictions_batch: PredictionsBatch):
+    activity_return = MagicMock()
+    activity_return.cleanup_prefix.return_value = None
+    workflow_mock.execute_activity_method.return_value = activity_return
+
     input_data = {
         'schedule_name': 'test_schedule',
         'model_name': 'test_model',
@@ -62,7 +65,8 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
     )
     prediction_input = {
         'metadata': metadata,
-        'data': {'success': True, 'data': {'col': ['test_data']}},
+        'data': activity_return,
+        'cleanup_prefix': activity_return.cleanup_prefix(),
         'schema': input_data['schema'],
         'table_name': input_data['table_name'],
         'transform_table_name': input_data['transform_table_name'],