sientia-dataops-model-manager/e2e/helpers.py

"""
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)