Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_lab_execution_postgres.py

@@ -0,0 +1,167 @@
+"""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()