sientia-dataops-laborious_temporal/e2e/helpers.py

"""
Shared helpers for E2E tests (Temporal workflows + PostgreSQL).
"""

import asyncio
from datetime import datetime
from decimal import Decimal
from typing import Any

from sqlalchemy import text
from sqlalchemy.engine import Engine


async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0):
    """
    Start a workflow and wait for its result.

    Args:
        client: Temporal client from WorkflowEnvironment.
        workflow_run: Workflow run method (e.g. PredictionsBatch.run).
        input_data: Workflow input payload.
        workflow_id: Unique workflow id.
        timeout: Max seconds to wait for completion.

    Return:
        Workflow result value.
    """
    handle = await client.start_workflow(
        workflow_run,
        input_data,
        id=workflow_id,
        task_queue='test-queue',
    )
    return await asyncio.wait_for(handle.result(), timeout=timeout)


def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> None:
    """
    Replace laborious_data rows for a model_id with one row per value (sensor_1..n).

    Args:
        postgres_engine: SQLAlchemy engine.
        model_id: Model id column value.
        values: Per-sensor values; use string 'NULL' for SQL NULL.
    """
    with postgres_engine.begin() as conn:
        conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}'))
        values_sql = []
        for i, value in enumerate(values):
            values_sql.append(f"""
                ({model_id}, 'sensor_{i + 1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
            """)
        insert_sql = f"""
            INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
            VALUES
                {', '.join(values_sql)}
        """
        conn.execute(text(insert_sql))


def assert_prediction(
    postgres_engine: Engine,
    model_id: int,
    prediction: float = 0.5,
    prediction_confidence: int | Decimal = 0,
    prediction_status: str = 'Good',
    comments: str = '',
) -> None:
    """
    Assert exactly one prediction row exists for model_id with expected columns.

    Args:
        postgres_engine: SQLAlchemy engine.
        model_id: Expected model_id.
        prediction: Expected prediction value.
        prediction_confidence: Expected confidence (int or Decimal for numeric column).
        prediction_status: Expected status string.
        comments: Expected comments string.
    """
    import pytest

    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(
                f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
                f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
                f'ORDER BY created_at ASC'
            )
        )
        prediction_rows = result_query.fetchall()
        assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}'
        row = prediction_rows[0]
        assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}'
        assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), (
            f'Expected prediction={prediction}, got {row[1]}'
        )
        assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal(
            str(prediction_confidence)
        ), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
        assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}"
        assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"


def assert_continue(
    postgres_engine: Engine,
    model_id: int,
    prediction_confidence: Decimal = Decimal(2),
    comments: str = 'Input data with bad quality',
) -> None:
    """Assert one default-style prediction row after CONTINUE gate path."""
    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(
                f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments '
                f'FROM predictions_schema.predictions WHERE model_id = {model_id}'
            )
        )
        prediction_rows = result_query.fetchall()
        assert len(prediction_rows) == 1, 'Expected one prediction record despite warnings'
        row = prediction_rows[0]
        assert row[1] == 0, f'Expected prediction=0, got {row[1]}'
        assert row[2] == prediction_confidence, (
            f'Expected prediction_confidence={prediction_confidence}, got {row[2]}'
        )
        assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}"
        assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"


def assert_stop(postgres_engine: Engine, model_id: int) -> None:
    """Assert no prediction rows for model_id."""
    import pytest

    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}')
        )
        count = result_query.scalar()
        assert count == 0, f'Expected no predictions, but found {count} records'


def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple) -> None:
    """
    Assert two prediction rows for model_id both match last_prediction.

    Rows are compared in created_at order for stability.

    Args:
        postgres_engine: SQLAlchemy engine.
        model_id: Model id.
        last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows.
    """
    import pytest

    with postgres_engine.connect() as conn:
        result_query = conn.execute(
            text(
                f'SELECT model_id, prediction, prediction_confidence, prediction_status '
                f'FROM predictions_schema.predictions WHERE model_id = {model_id} '
                f'ORDER BY created_at ASC'
            )
        )
        prediction_rows = result_query.fetchall()
        assert len(prediction_rows) == 2, 'Expected two prediction records'
        assert prediction_rows[0] == last_prediction, (
            f'Expected first row {last_prediction}, got {prediction_rows[0]}'
        )
        assert prediction_rows[1] == last_prediction, (
            f'Expected second row {last_prediction}, got {prediction_rows[1]}'
        )


def make_workflow_id(prefix: str) -> str:
    """Build a unique workflow id using a prefix and current timestamp."""
    return f'{prefix}-{datetime.now().timestamp()}'