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

"""P3 checkpoint test: search import_state + AdaptiveLoop export/restore.

Verifies the /resume path:
  1. FeasibilityFirstSearch.export_state -> import_state round-trip keeps
     run_id / budget / points / objective / convergence and can select the
     next batch.
  2. AdaptiveLoop.export_state -> restore_state survives a simulated process
     restart (registry cleared) and can keep driving batches.

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

Isolated temp SQLite DB; no real Motor-CAD involved.
"""
import os
import sys
import tempfile

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

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

from app.services.feasibility_search import (  # noqa: E402
    FeasibilityFirstSearch, ParameterRange, L0PreScreeningEngine,
)

# ---------------------------------------------------------------- 1. search round-trip
params = [ParameterRange(name="airgap_mm", min_value=0.8, max_value=2.0,
                         step=0.1, unit="mm")]
s1 = FeasibilityFirstSearch(parameters=params, l0_engine=L0PreScreeningEngine(),
                            total_budget=8, batch_size=4, initial_samples=4,
                            objective_metric="tavg_nm", objective_direction="maximize")
batch = s1.generate_initial_batch()
for i, p in enumerate(batch[:3]):
    s1.report_result(p.id, {"tavg_nm": 8.0 + 0.5 * p.id}, "ok")

ex = s1.export_state()
s2 = FeasibilityFirstSearch.import_state(ex, l0_engine=L0PreScreeningEngine())
assert s2.state.run_id == s1.state.run_id
assert s2.state.used_budget == s1.state.used_budget == 3
assert len(s2.state.points) == len(s1.state.points)
assert s2.objective_metric == "tavg_nm" and s2.objective_direction == "maximize"
assert s2.state.convergence_status == s1.state.convergence_status
nb = s2.select_next_batch()
assert nb, "restored search must be able to select the next batch"
print("[1] search export->import round-trip OK (next batch: %d pts)" % len(nb))

# ---------------------------------------------------------------- 2. AdaptiveLoop resume
from app.services.adaptive_loop import (  # noqa: E402
    AdaptiveLoop, create_loop, _loops,
)

loop = create_loop(user_requirement="max torque", total_budget=8, batch_size=4)
loop.plan = {
    "plan_name": "p", "topology": "SSSR",
    "scan_variables": [
        {"name": "airgap_mm", "min_value": 0.8, "max_value": 2.0},
        {"name": "current_a", "min_value": 5.0, "max_value": 20.0},
    ],
    "search_strategy": {"max_solver_calls": 8, "batch_size": 4,
                        "initial_samples": 4},
    "acceptance_criteria": {"objective_metric": "tavg_nm",
                            "objective_direction": "maximize",
                            "hard_constraints": []},
}
loop.initialize_search()
sub = loop.submit_batch_to_executor()
assert sub.get("task_id"), sub
exported = loop.export_state()
assert exported["search"] is not None
assert exported["loop_id"] == loop.loop_id
assert exported["phase"] == loop.phase.value

# simulate process restart: in-memory registry is lost
_loops.clear()
assert len(_loops) == 0

loop2 = AdaptiveLoop.restore_state(exported)
assert loop2.loop_id == loop.loop_id
assert loop2.phase.value == exported["phase"]
assert loop2.search is not None
sub2 = loop2.submit_batch_to_executor()
assert sub2.get("task_id"), "restored loop must be able to submit a batch"
print("[2] AdaptiveLoop export->restore after restart OK (batch resubmitted)")

print("\nALL P3 CHECKPOINT TESTS PASSED")