Code import - branch release/SIENTIAPDE-1645

e2e/helpers.py

@@ -0,0 +1,195 @@
+"""
+Shared helpers for E2E tests (Temporal workflows + PostgreSQL).
+"""
+
+import asyncio
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import pytest
+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 = 600.0,
+):
+    """
+    Start a Temporal workflow and wait for its result.
+
+    Args:
+        client: Temporal client from WorkflowEnvironment.
+        workflow_run: Workflow run method (e.g. TrainModel.run).
+        input_data: Workflow input payload.
+        workflow_id: Unique workflow id.
+        timeout: Max seconds to wait for completion (default allows cold testcontainer startup).
+
+    Returns:
+        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 make_workflow_id(prefix: str) -> str:
+    """Build a unique workflow id using a prefix and current timestamp."""
+    return f'{prefix}-{datetime.now().timestamp()}'
+
+
+def insert_experiment_run(
+    engine: Engine,
+    experiment_run_id: int,
+    experiment_name: str = 'test_experiment',
+    status: str = 'ORCHESTRATOR_WAITING_PROC',
+    bucket_name: str = 'model-training',
+    file_name: str = 'training_data.csv',
+) -> None:
+    """
+    Insert a minimal experiment_run row to satisfy foreign-key-style lookups.
+
+    Args:
+        engine: SQLAlchemy engine connected to the test database.
+        experiment_run_id: Primary key for the row.
+        experiment_name: Human-readable experiment name.
+        status: Initial status string.
+        bucket_name: MinIO bucket name.
+        file_name: Training file name inside the bucket.
+    """
+    with engine.begin() as conn:
+        conn.execute(
+            text("""
+                INSERT INTO public.experiment_run
+                    (id, experiment_name, status, bucket_name, file_name)
+                VALUES
+                    (:id, :experiment_name, :status, :bucket_name, :file_name)
+                ON CONFLICT (id) DO NOTHING
+            """),
+            {
+                'id': experiment_run_id,
+                'experiment_name': experiment_name,
+                'status': status,
+                'bucket_name': bucket_name,
+                'file_name': file_name,
+            },
+        )
+
+
+def assert_experiment_status(
+    engine: Engine,
+    experiment_run_id: int,
+    expected_status: str,
+) -> None:
+    """
+    Assert the final status of an experiment_run row.
+
+    Args:
+        engine: SQLAlchemy engine.
+        experiment_run_id: Row primary key.
+        expected_status: Expected status string.
+    """
+    with engine.connect() as conn:
+        row = conn.execute(
+            text('SELECT status FROM public.experiment_run WHERE id = :id'),
+            {'id': experiment_run_id},
+        ).fetchone()
+
+    assert row is not None, (
+        f'No experiment_run row found for id={experiment_run_id}'
+    )
+    assert row[0] == expected_status, (
+        f'Expected status={expected_status!r}, got {row[0]!r} '
+        f'for experiment_run id={experiment_run_id}'
+    )
+
+
+def assert_experiment_run_name_set(
+    engine: Engine,
+    experiment_run_id: int,
+) -> None:
+    """Assert that run_name is not null/empty after a successful training."""
+    with engine.connect() as conn:
+        row = conn.execute(
+            text('SELECT run_name FROM public.experiment_run WHERE id = :id'),
+            {'id': experiment_run_id},
+        ).fetchone()
+
+    assert row is not None, (
+        f'No experiment_run row found for id={experiment_run_id}'
+    )
+    assert row[0] is not None and row[0].strip() != '', (
+        f'Expected run_name to be set for experiment_run id={experiment_run_id}, got {row[0]!r}'
+    )
+
+
+def assert_experiment_error(
+    engine: Engine,
+    experiment_run_id: int,
+    expected_status: str,
+    error_substr: str,
+) -> None:
+    """
+    Assert status and that error_message contains a given substring.
+
+    Args:
+        engine: SQLAlchemy engine.
+        experiment_run_id: Row primary key.
+        expected_status: Expected status string.
+        error_substr: Substring that must appear in error_message.
+    """
+    with engine.connect() as conn:
+        row = conn.execute(
+            text(
+                'SELECT status, error_message FROM public.experiment_run WHERE id = :id'
+            ),
+            {'id': experiment_run_id},
+        ).fetchone()
+
+    assert row is not None, (
+        f'No experiment_run row found for id={experiment_run_id}'
+    )
+    assert row[0] == expected_status, (
+        f'Expected status={expected_status!r}, got {row[0]!r}'
+    )
+    assert row[1] is not None and error_substr.lower() in row[1].lower(), (
+        f'Expected error_message to contain {error_substr!r}, got {row[1]!r}'
+    )
+
+
+def assert_no_experiment_row(engine: Engine, experiment_run_id: int) -> None:
+    """Assert that no experiment_run row exists for the given id."""
+    with engine.connect() as conn:
+        count = conn.execute(
+            text('SELECT COUNT(*) FROM public.experiment_run WHERE id = :id'),
+            {'id': experiment_run_id},
+        ).scalar()
+    assert count == 0, (
+        f'Expected no experiment_run row for id={experiment_run_id}, found {count}'
+    )
+
+
+def load_scenario(scenario_filename: str) -> dict[str, Any]:
+    """
+    Load a test scenario JSON file from docs/test-scenarios/.
+
+    Args:
+        scenario_filename: Filename without path (e.g. '01-linear-regression-basic.json').
+
+    Returns:
+        dict: Parsed scenario payload.
+    """
+    scenario_path = (
+        Path(__file__).parent.parent / 'docs' / 'test-scenarios' / scenario_filename
+    )
+    with open(scenario_path) as f:
+        return json.load(f)