Snapshot of fix/QTZPOC-13 source tree

Code-only import without upstream history.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

e2e/helpers.py

@@ -0,0 +1,186 @@
+"""
+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 = None,
+    comments_contains: str | None = None,
+) -> 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 exact comments string (optional).
+        comments_contains: Substring expected in comments when queued (optional).
+    """
+
+    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]}"
+        )
+        if comments is not None:
+            assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}"
+        if comments_contains is not None:
+            assert comments_contains in row[4], (
+                f"Expected comments to contain '{comments_contains}', 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."""
+
+    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.
+    """
+
+    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()}'