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

"""P3-M2 regression: AdaptiveOrchestrator full closed loop with fake executor.

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

Uses an isolated temp SQLite DB and temp loop-state dir, so it never touches
the real web DB or real loop state. No real Motor-CAD is involved.
"""
import json
import os
import sys
import tempfile

_TMP = os.path.join(tempfile.mkdtemp(), "test_afm.db")
_TMP_STATE = os.path.join(tempfile.mkdtemp(), "loops")
os.environ["AFM_DB_PATH"] = _TMP
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.strategy_orchestrator import AdaptiveOrchestrator  # noqa: E402
from app.services.task_manager import get_task_manager  # noqa: E402

tm = get_task_manager()
orch = AdaptiveOrchestrator(state_dir=_TMP_STATE)

# ------------------------------------------------------------------ 1. start
res = orch.start_loop(
    loop_id="loop-p3-test",
    parameters=[
        {"name": "airgap_mm", "min_value": 0.8, "max_value": 2.0, "step": 0.1, "unit": "mm"},
        {"name": "current_a", "min_value": 5.0, "max_value": 20.0, "step": 0.5},
    ],
    total_budget=8,
    batch_size=4,
    initial_samples=4,
    objective_metric="tavg_nm",
    objective_direction="maximize",
)
assert res["phase"] == "running", res
assert res["current_task_id"], res
assert res["batch_task"]["task_type"] == "adaptive_batch", res
first_batch_ids = res["batch_task"]["point_ids"]
assert isinstance(first_batch_ids, list) and len(first_batch_ids) > 0, first_batch_ids
print("[1] start_loop OK: task=%s first_batch=%d points"
      % (res["current_task_id"], len(first_batch_ids)))


def run_batch(tid):
    """Fake executor: read task.json parameters, fabricate results, report."""
    task = tm.get_task(tid)
    with open(task["task_file"], "r", encoding="utf-8") as f:
        payload = json.load(f)
    results = []
    for i, params in enumerate(payload["parameters"]):
        pid = params.get("point_id")
        results.append({
            "point_id": pid,
            "point_index": i,
            "params": params,
            "metrics": {"tavg_nm": round(8.0 + 0.5 * pid, 3), "efficiency_pct": 90.0 + (pid % 5)},
            "status": "OK",
        })
    tm.report_results(tid, results, status="completed")
    return results


# ---------------------------------------------------------------- 2. drive
steps = 0
max_steps = 12
batches_seen = set()
while steps < max_steps:
    steps += 1
    view = orch.get_loop_status("loop-p3-test")
    if view["phase"] in ("converged", "budget_exhausted", "failed"):
        break
    tid = view["current_task_id"]
    if tid is None:
        view = orch.advance_loop("loop-p3-test")
        continue
    run_batch(tid)
    view = orch.advance_loop("loop-p3-test")
    batches_seen.add(view.get("current_batch"))
    print("[2] step %d -> batch=%s task=%s phase=%s n_results=%s"
          % (steps, view.get("current_batch"), view.get("current_task_id"),
             view["phase"], view["n_results"]))

final = orch.get_loop_status("loop-p3-test")
print("[3] FINAL phase=%s batches=%s n_results=%s"
      % (final["phase"], sorted(batches_seen), final["n_results"]))
assert final["phase"] in ("converged", "budget_exhausted"), final
assert final["n_results"] >= 4, final
assert len(batches_seen) >= 1, batches_seen
# all reported points must be reflected in search state
search_state = final.get("search_state") or {}
assert search_state.get("completed_points", 0) >= 4, search_state
assert search_state.get("used_budget", 0) > 0, search_state

# ------------------------------------------------------------- 3. misc
loops = orch.list_loops()
assert len(loops) == 1 and loops[0]["loop_id"] == "loop-p3-test", loops

try:
    orch.start_loop(loop_id="loop-p3-test",
                    parameters=[{"name": "airgap_mm", "min_value": 1, "max_value": 2}])
    raise SystemExit("duplicate loop should fail")
except ValueError:
    print("[4] duplicate loop rejected OK")

state_path = os.path.join(_TMP_STATE, "loop-p3-test_loop.json")
assert os.path.exists(state_path), state_path
print("[5] loop state persisted OK")

# ------------------------------------------------------------ 4. task model
# create_task with explicit adaptive-batch fields
t2 = tm.create_task(
    plan_id=None,
    plan_data={"topology": "SSSR"},
    parameters=[{"airgap_mm": 1.0, "point_id": 0}],
    task_name="meta-batch",
    task_type="adaptive_batch",
    loop_id="loop-x",
    batch_id=3,
    point_ids=[0, 1],
    dynamic=True,
)
assert t2["task_type"] == "adaptive_batch", t2
assert t2["loop_id"] == "loop-x" and t2["batch_id"] == 3, t2
assert t2["point_ids"] == [0, 1] and t2["dynamic"] is True, t2
print("[6] task model adaptive-batch fields OK")

print("\nALL P3-M2 ORCHESTRATOR TESTS PASSED")