Code import - branch release/SIENTIAPDE-1646

e2e/test_minio_offload.py

@@ -0,0 +1,108 @@
+"""
+E2E-style tests for MinIO offload using a real MinIO testcontainer.
+"""
+
+from unittest.mock import patch
+
+import pytest
+from sqlalchemy import text
+from temporalio.testing import WorkflowEnvironment
+from temporalio.worker import Worker
+
+from e2e.helpers import (
+    insert_sample_data,
+    load_scenario_input,
+    make_workflow_id,
+    start_and_await_workflow,
+)
+from laborious.activities.activities import Activities
+from laborious.utils.models import minio_dataframe_payload as mdp
+from laborious.workflows.predictions_batch import PredictionsBatch
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_load_query_with_minio_offload_writes_object_to_bucket(
+    postgres_engine,
+    minio_container,
+    test_activities_real_minio: Activities,
+):
+    """
+    With a tiny offload threshold, query results are uploaded as Parquet to MinIO.
+
+    Uses real MinioRepository against testcontainers MinIO (no MinIO mock).
+    """
+    model_id = 501
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
+    insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
+
+    scenario_input = load_scenario_input('minio_offload_load_query.json', model_id=model_id)
+    metadata = {'metadata': scenario_input['metadata']}
+    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
+        payload = test_activities_real_minio.load_query_with_minio_offload(scenario_input)
+    assert payload.object_key, 'offloaded payload must reference a MinIO object'
+    assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict'
+
+    df = payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata'])
+    assert len(df) >= 1
+
+    client = minio_container.get_client()
+    listed = list(client.list_objects('test-bucket', recursive=True))
+    names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed]
+    assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}'
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_predictions_batch_with_minio_offload_path(
+    temporal_test_env: WorkflowEnvironment,
+    temporal_worker_real_minio: Worker,
+    postgres_engine,
+    test_activities_real_minio: Activities,
+):
+    """
+    Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes.
+    """
+    model_id = 502
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
+        conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}'))
+        conn.execute(text(f'DELETE FROM sientia_data.transformed_data WHERE model_id = {model_id}'))
+    insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])
+
+    input_data = load_scenario_input('minio_offload_workflow.json', model_id=model_id)
+
+    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
+        await start_and_await_workflow(
+            temporal_test_env.client,
+            PredictionsBatch.run,
+            input_data,
+            make_workflow_id('test-batch-minio-offload'),
+        )
+
+    with postgres_engine.connect() as conn:
+        count = conn.execute(
+            text(f'SELECT COUNT(*) FROM sientia_data.predictions WHERE model_id = {model_id}')
+        ).scalar()
+    assert count == 1
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_load_query_with_inline_payload_when_below_threshold(
+    postgres_engine,
+    test_activities_real_minio: Activities,
+):
+    """Scenario 4.2.1: payload stays inline when threshold is high enough."""
+    model_id = 503
+    with postgres_engine.begin() as conn:
+        conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}'))
+    insert_sample_data(postgres_engine, model_id, [1.0, 2.0])
+
+    scenario_input = load_scenario_input('minio_offload_load_query.json', model_id=model_id)
+    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 10**9):
+        payload = test_activities_real_minio.load_query_with_minio_offload(scenario_input)
+
+    assert payload.object_key is None
+    assert payload.data is not None