Publish curated ITWorx-ModelForge source from 990401737f96c8d46a9f7b0d84994516cadcb6a0

backend/tests/test_concurrency_m16.py

@@ -12,6 +12,7 @@ the live chaos evidence; what these tests pin down is that the *code* claims a s
 
 from __future__ import annotations
 
+import importlib
 import os
 import threading
 import uuid
@@ -22,9 +23,10 @@ from pathlib import Path
 from typing import Any
 
 import pytest
+from alembic.config import Config
 from pydantic import SecretStr
 from sqlalchemy import create_engine, event, func, select, text
-from sqlalchemy.engine import Engine
+from sqlalchemy.engine import Engine, make_url
 from sqlalchemy.exc import IntegrityError, OperationalError
 from sqlalchemy.orm import Session
 
@@ -151,7 +153,9 @@ def test_a_single_use_enrolment_token_survives_a_concurrent_storm(
     with Session(engine) as session:
         assert session.scalar(select(func.count()).select_from(ComputeNode)) == 1
         active = session.scalar(
-            select(func.count()).select_from(NodeCredential).where(NodeCredential.revoked_at.is_(None))
+            select(func.count())
+            .select_from(NodeCredential)
+            .where(NodeCredential.revoked_at.is_(None))
         )
         assert active == 1
         enrollment = session.scalar(select(NodeEnrollment))
@@ -168,7 +172,7 @@ def test_managed_postgresql_claim_and_revocation_linearize_on_the_row_lock() ->
     """Managed gate for the PostgreSQL row-lock semantics SQLite cannot reproduce.
 
     Set MODELFORGE_TEST_POSTGRES_URL to an isolated, disposable PostgreSQL database. The test
-    creates and removes one UUID-named schema, pauses enrolment after its conditional claim has
+    creates and removes one UUID-named database, pauses enrolment after its conditional claim has
     acquired the row lock, then proves a concurrent revocation loses after the claim commits.
     """
 
@@ -180,15 +184,17 @@ def test_managed_postgresql_claim_and_revocation_linearize_on_the_row_lock() ->
     if database_url.startswith("postgresql://"):
         database_url = database_url.replace("postgresql://", "postgresql+psycopg://", 1)
 
-    schema = f"test_enrollment_linearization_{uuid.uuid4().hex}"
-    admin_engine = create_engine(database_url, pool_pre_ping=True)
-    with admin_engine.begin() as connection:
-        connection.execute(text(f'CREATE SCHEMA "{schema}"'))
-    engine = create_engine(
-        database_url,
-        pool_pre_ping=True,
-        connect_args={"options": f"-csearch_path={schema}"},
-    )
+    owner_url = os.getenv("MODELFORGE_TEST_POSTGRES_OWNER_URL")
+    runtime_url = os.getenv("MODELFORGE_TEST_POSTGRES_RUNTIME_URL")
+    if not owner_url or not runtime_url:
+        pytest.fail("managed enrollment test requires owner and runtime URLs")
+    database = f"test_enrollment_linearization_{uuid.uuid4().hex}"
+    admin_engine = create_engine(database_url, isolation_level="AUTOCOMMIT")
+    with admin_engine.connect() as connection:
+        connection.execute(text(f'CREATE DATABASE "{database}" OWNER modelforge'))
+    # The canonical audit function deliberately targets public tables. A search_path-only
+    # fixture would read a private checkpoint but append to another schema's audit chain.
+    engine = create_engine(make_url(runtime_url).set(database=database), pool_pre_ping=True)
     claim_written = threading.Event()
     allow_claim_commit = threading.Event()
     revoke_update_started = threading.Event()
@@ -207,7 +213,18 @@ def test_managed_postgresql_claim_and_revocation_linearize_on_the_row_lock() ->
             revoke_update_started.set()
 
     try:
-        Base.metadata.create_all(engine)
+        migration = Config()
+        migration.set_main_option(
+            "script_location", str(Path(__file__).resolve().parents[1] / "alembic")
+        )
+        migration.set_main_option(
+            "sqlalchemy.url",
+            make_url(owner_url)
+            .set(database=database)
+            .render_as_string(hide_password=False)
+            .replace("%", "%%"),
+        )
+        importlib.import_module("alembic.command").upgrade(migration, "head")
         with Session(engine) as session:
             created = NodeAgentService(session, settings()).create_enrollment(
                 EnrollmentTokenCreate(display_name="PostgreSQL row-lock target")
@@ -293,8 +310,8 @@ def test_managed_postgresql_claim_and_revocation_linearize_on_the_row_lock() ->
     finally:
         allow_claim_commit.set()
         engine.dispose()
-        with admin_engine.begin() as connection:
-            connection.execute(text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
+        with admin_engine.connect() as connection:
+            connection.execute(text(f'DROP DATABASE IF EXISTS "{database}"'))
         admin_engine.dispose()
 
 
@@ -360,9 +377,7 @@ def test_re_enrolment_of_one_identity_under_load_keeps_one_active_credential(
     with Session(engine) as session:
         service = NodeAgentService(session, settings())
         tokens = [
-            service.create_enrollment(
-                EnrollmentTokenCreate(display_name="GPU Node")
-            ).enrollment_token
+            service.create_enrollment(EnrollmentTokenCreate(display_name="GPU Node")).enrollment_token
             for _ in range(rounds)
         ]
 
@@ -421,9 +436,7 @@ def test_the_same_idempotency_key_in_parallel_produces_one_logical_operation(
 
     def attempt(_index: int) -> str:
         with Session(engine) as session:
-            existing = session.scalar(
-                select(ServingJob).where(ServingJob.idempotency_key == key)
-            )
+            existing = session.scalar(select(ServingJob).where(ServingJob.idempotency_key == key))
             if existing is not None:
                 return "REUSED"
             session.add(
@@ -478,9 +491,7 @@ def test_concurrent_audit_appends_form_one_strict_chain(engine: Engine) -> None:
         assert len(events) == attempts
         assert audit_chain_violations(events, session.get(AuditChainHead, 1)) == []
         result = next(
-            item
-            for item in check_invariants(session).results
-            if item.key == "audit_chain_intact"
+            item for item in check_invariants(session).results if item.key == "audit_chain_intact"
         )
         assert result.status is InvariantStatus.HOLDS