sientia-dataops-laborious_temporal/e2e/test_minio_offload.py

"""
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, 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 predictions_schema.laborious_data WHERE model_id = {model_id}')
        )
    insert_sample_data(postgres_engine, model_id, [1.0, 2.0])

    metadata = {
        'metadata': {
            'schedule_name': 'test-schedule',
            'model_name': 'test_model',
            'model_id': model_id,
            'workflow_name': 'predictions_batch',
        }
    }
    with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1):
        payload = await test_activities_real_minio.load_query_with_minio_offload(
            {
                **metadata,
                'query': (
                    'SELECT timestamp, variable, value, created_at '
                    f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
                ),
                'model_name': 'test_model',
                'datetime_columns': ['timestamp', 'created_at'],
            }
        )
    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 = await 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 predictions_schema.laborious_data WHERE model_id = {model_id}')
        )
        conn.execute(
            text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')
        )
        conn.execute(
            text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')
        )
    insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0])

    input_data = {
        'schedule_name': 'test-schedule',
        'model_name': 'test_model',
        'model_id': model_id,
        'query': (
            'SELECT timestamp, variable, value, created_at '
            f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}'
        ),
        'schema': 'predictions_schema',
        'table_name': 'predictions',
        'transform_table_name': 'transformed_data',
        'input_filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}},
        'mlflow_transform_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
        'mlflow_predict_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}},
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
        'opc_output_config': {},
        'pi_web_api_output_config': {},
        'save_transform': True,
        'prediction_store_policy': 'lts:1',
        'model_config': {
            'retention_minutes': 0,
            'transform_flavor': 'sklearn',
            'predict_flavor': 'sklearn',
        },
        'datetime_columns': ['timestamp', 'created_at'],
    }

    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 predictions_schema.predictions WHERE model_id = {model_id}')
        ).scalar()
    assert count == 1