ModelForge-Public/backend/tests/test_lab_execution_postgres.py

"""Managed PostgreSQL migration and actor-scoped LAB idempotency qualification."""

from __future__ import annotations

import importlib
import os
import threading
import uuid
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from types import SimpleNamespace

import pytest
from alembic.config import Config
from sqlalchemy import create_engine, func, select, text
from sqlalchemy.engine import make_url
from sqlalchemy.orm import Session

from modelforge_api.domain.lab import LabExecutionStartRequest
from modelforge_api.persistence.models import GatewayRequest, LabExecutionRun
from modelforge_api.services.lab_runs import LabRunService
from modelforge_api.services.serving import ServingService
from modelforge_api.services.transient_payloads import MemoryPayloadStore
from tests.test_dispatch_context import operator_fixture

_PRIOR = "20260923_0037"
_CURRENT = "20260924_0039"


def test_managed_postgresql_lab_migration_and_concurrent_idempotent_start(
) -> None:
    """The database constraint must linearize simultaneous starts for one actor/key."""
    admin_url = os.getenv("MODELFORGE_TEST_POSTGRES_URL")
    if not admin_url:
        pytest.skip("requires managed MODELFORGE_TEST_POSTGRES_URL")
    owner_url = os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL")
    runtime_url = os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL")
    assert owner_url and runtime_url, "managed qualification requires owner and runtime URLs"

    if admin_url.startswith("postgresql://"):
        admin_url = admin_url.replace("postgresql://", "postgresql+psycopg://", 1)
    owner = make_url(owner_url)
    runtime = make_url(runtime_url)
    assert owner.get_backend_name() == runtime.get_backend_name() == "postgresql"
    assert owner.username == "modelforge"
    assert runtime.username == "modelforge_runtime"

    database_name = f"test_lab_runs_{uuid.uuid4().hex}"
    admin = create_engine(admin_url, isolation_level="AUTOCOMMIT")
    owner_engine = create_engine(owner.set(database=database_name))
    runtime_engine = create_engine(
        runtime.set(database=database_name),
        connect_args={"connect_timeout": 15, "options": "-c statement_timeout=30000 -c lock_timeout=15000"},
    )
    migration = Config()
    migration.set_main_option(
        "script_location", str(Path(__file__).resolve().parents[1] / "alembic")
    )
    migration.set_main_option(
        "sqlalchemy.url",
        owner.set(database=database_name)
        .render_as_string(hide_password=False)
        .replace("%", "%%"),
    )
    alembic = importlib.import_module("alembic.command")
    barrier = threading.Barrier(2)

    with admin.connect() as connection:
        connection.execute(text(f'CREATE DATABASE "{database_name}" OWNER modelforge'))
    try:
        alembic.upgrade(migration, _PRIOR)
        with runtime_engine.connect() as connection:
            assert connection.scalar(text("SELECT to_regclass('public.lab_execution_runs')")) is None

        alembic.upgrade(migration, "head")
        with runtime_engine.connect() as connection:
            assert connection.scalar(text("SELECT version_num FROM alembic_version")) == _CURRENT
            assert connection.scalar(
                text("SELECT has_table_privilege(current_user, 'lab_execution_runs', 'SELECT')")
            )
            for privilege in ("INSERT", "UPDATE"):
                assert connection.scalar(
                    text(
                        "SELECT has_table_privilege(current_user, 'lab_execution_runs', :privilege)"
                    ),
                    {"privilege": privilege},
                )
            for privilege in ("DELETE", "TRUNCATE"):
                assert not connection.scalar(
                    text(
                        "SELECT has_table_privilege(current_user, 'lab_execution_runs', :privilege)"
                    ),
                    {"privilege": privilege},
                )

            with Session(owner_engine) as setup_session:
                setup_serving, preflight, selected = operator_fixture(setup_session)
                request = LabExecutionStartRequest(
                    target=preflight.target,
                    input=preflight.input,
                    limits=preflight.limits,
                    deployment_id=selected.capability_deployment_id,
                    experiment_route_id=selected.id,
                )
                deployment_settings = setup_serving._deployment_settings
                manifests = setup_serving.manifests

        def start_once() -> uuid.UUID:
            with Session(runtime_engine) as session:
                serving = ServingService(
                    session,
                    deployment_settings,
                    manifests,
                    MemoryPayloadStore(),
                )
                service = LabRunService(serving)
                select_target = service._selected_target

                def select_same_target(_request: LabExecutionStartRequest):
                    selection = select_target(_request)
                    barrier.wait(timeout=30)
                    return selection

                service._selected_target = select_same_target  # type: ignore[method-assign]
                result = service.start(
                    request,
                    actor_subject="postgres-race-operator",
                    idempotency_key="same-selected-run",
                )
                return result.id

        with ThreadPoolExecutor(max_workers=2) as pool:
            run_ids = list(pool.map(lambda _index: start_once(), range(2)))
        assert run_ids[0] == run_ids[1]

        cancel_barrier = threading.Barrier(2)

        def cancel_once() -> str:
            with Session(runtime_engine) as session:
                serving = SimpleNamespace(session=session, payload_store=MemoryPayloadStore())
                service = LabRunService(serving)
                cancel_barrier.wait(timeout=30)
                return service.cancel(
                    run_ids[0], actor_subject="postgres-race-operator"
                ).status

        with ThreadPoolExecutor(max_workers=2) as pool:
            cancel_statuses = list(pool.map(lambda _index: cancel_once(), range(2)))
        assert cancel_statuses == ["cancelled", "cancelled"]

        with Session(runtime_engine) as session:
            assert session.scalar(select(func.count()).select_from(LabExecutionRun)) == 1
            assert session.scalar(select(func.count()).select_from(GatewayRequest)) == 1
            run = session.scalar(select(LabExecutionRun))
            assert run is not None
            assert run.actor_subject == "postgres-race-operator"
            assert run.idempotency_key == "same-selected-run"
            assert run.status == "cancelled"
            assert run.cancellation_requested_at is not None
            assert run.cancellation_acknowledged_at is None
            assert run.input_sha256 and "private postgres race input" not in str(run.target_snapshot)
    finally:
        runtime_engine.dispose()
        owner_engine.dispose()
        with admin.connect() as connection:
            connection.execute(text(f'DROP DATABASE IF EXISTS "{database_name}" WITH (FORCE)'))
        admin.dispose()