pcb-afm-simulation-system/scripts/test_p3_concurrency.py

"""P3 concurrency test: atomic task claim under parallel executor instances.

Verifies that when N executors race to claim the same pending task, exactly
one wins; the rest get ValueError (already claimed) or a transient
SQLite lock error - never a duplicate successful claim (which would cause
double simulation of the same points).

Run:  python scripts/test_p3_concurrency.py   (exit 0 = PASS)

Isolated temp SQLite DB; no real Motor-CAD involved.
"""
import os
import sys
import tempfile
import threading
from concurrent.futures import ThreadPoolExecutor

_TMP = os.path.join(tempfile.mkdtemp(), "conc.db")
os.environ["AFM_DB_PATH"] = _TMP
os.environ["KIMI_API_KEY"] = ""
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "web", "backend"))
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.database import init_db  # noqa: E402
init_db()

from app.services.task_manager import get_task_manager  # noqa: E402

tm = get_task_manager()


def make_pending_task():
    t = tm.create_task(
        plan_id=None, plan_data={},
        parameters=[{"airgap_mm": 1.0, "point_id": 1}],
        task_name="race-task", task_type="scan",
    )
    return t["task_id"]


# ------------------------------------------------------------------ 1. manager-level race
N = 8
tid = make_pending_task()
success = []
failures = []
lock = threading.Lock()


def claim():
    try:
        tm.dispatch_task(tid)
        with lock:
            success.append(1)
    except Exception as exc:  # ValueError or transient lock error
        with lock:
            failures.append(type(exc).__name__)


with ThreadPoolExecutor(max_workers=N) as pool:
    list(pool.map(lambda _: claim(), range(N)))

assert len(success) == 1, "exactly one claim must win, got %d" % len(success)
assert len(failures) == N - 1, "the rest must fail, got %d failures" % len(failures)
print("[1] manager-level race: 1 win / %d lost (%s)" % (len(failures), set(failures)))

# final state must be dispatched
task = tm.get_task(tid)
assert task["status"] == "dispatched", task
print("[2] final status = dispatched OK")

# a sequential second claim must be rejected
try:
    tm.dispatch_task(tid)
    raise SystemExit("second claim should fail")
except ValueError:
    print("[3] sequential second claim rejected OK")

# ------------------------------------------------------------------ 2. HTTP-level race
from app.database import SessionLocal  # noqa: E402
from app.models.task import Task  # noqa: E402

# create a second pending task directly for the HTTP race
with SessionLocal() as db:
    task = Task(
        task_id="http-race-1",
        task_name="http-race", plan_id=None,
        status="pending", priority=5, created_by="test",
    )
    db.add(task)
    db.commit()

http_success = []


def http_claim():
    try:
        tm.dispatch_task("http-race-1")
        with lock:
            http_success.append(1)
    except Exception:
        pass


with ThreadPoolExecutor(max_workers=N) as pool:
    list(pool.map(lambda _: http_claim(), range(N)))
assert len(http_success) == 1, "HTTP-level exactly-one claim, got %d" % len(http_success)
print("[4] HTTP-level race: exactly 1 claim won OK")

print("\nALL P3 CONCURRENCY TESTS PASSED")