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

"""Tests for plan_schema module - especially string enum parameter support.

Covers:
- FixedParam with string enum values (Winding_Connection, Cooling_Type, etc.)
- Validation of plans with mixed numeric/string params
- generate_full_params preserves string values
- Boundary cases (empty, int, float, str values)

Run: python scripts/test_plan_schema.py
Exit 0 = PASS, non-zero = FAIL.
"""
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "web", "backend"))

from src.plan_schema import (
    FixedParam,
    ScanVariable,
    SimulationPlan,
    validate_plan_dict,
    parse_plan,
)


def _make_plan(fixed_params, variables=None):
    """Helper to create a minimal valid plan dict."""
    return {
        "plan_id": "test-plan",
        "plan_version": "2.0",
        "created_at": "2026-08-30T12:00:00",
        "topology": "SSSR",
        "model_path": "models/test.mot",
        "fixed_params": fixed_params,
        "variables": variables or [
            {"name": "Magnet_Thickness", "values": [2, 3], "unit": "mm"},
        ],
        "cases": [{"id": "default", "name": "Default", "params": {}}],
        "output_metrics": ["tavg_nm", "efficiency_pct"],
    }


def test_fixed_param_string_enum():
    """Test that FixedParam accepts string enum values.

    Regression test for the AI generate-and-save 422 error:
    Winding_Connection="Star" caused 'could not convert string to float'.
    """
    fp = FixedParam.from_dict({
        "name": "Winding_Connection",
        "value": "Star",
        "unit": "",
    })
    assert fp.value == "Star"
    assert isinstance(fp.value, str)
    print("PASS: test_fixed_param_string_enum")


def test_fixed_param_numeric_values():
    """Test that numeric values still work correctly."""
    fp_int = FixedParam.from_dict({"name": "Slot_Number", "value": 12})
    assert fp_int.value == 12
    assert isinstance(fp_int.value, int)

    fp_float = FixedParam.from_dict({"name": "Airgap", "value": 1.5})
    assert fp_float.value == 1.5
    assert isinstance(fp_float.value, float)
    print("PASS: test_fixed_param_numeric_values")


def test_fixed_param_to_dict_roundtrip():
    """Test that string values survive to_dict -> from_dict roundtrip."""
    original = FixedParam(name="Cooling_Type", value="Natural Convection", unit="")
    d = original.to_dict()
    assert d["value"] == "Natural Convection"
    restored = FixedParam.from_dict(d)
    assert restored.value == "Natural Convection"
    assert isinstance(restored.value, str)
    print("PASS: test_fixed_param_to_dict_roundtrip")


def test_validate_plan_with_string_enums():
    """Test that a plan with mixed numeric/string fixed params validates OK."""
    plan = _make_plan([
        {"name": "Airgap", "value": 1.0, "unit": "mm"},
        {"name": "Winding_Connection", "value": "Star", "unit": ""},
        {"name": "Cooling_Type", "value": "Natural Convection", "unit": ""},
        {"name": "Slot_Number", "value": 12, "unit": ""},
    ])
    ok, errors = validate_plan_dict(plan, require_model_path=False)
    assert ok, f"Plan should validate OK, got errors: {errors}"
    assert len(errors) == 0
    print("PASS: test_validate_plan_with_string_enums")


def test_generate_full_params_preserves_strings():
    """Test that generate_full_params preserves string enum values in all points."""
    plan = _make_plan([
        {"name": "Airgap", "value": 1.0, "unit": "mm"},
        {"name": "Winding_Connection", "value": "Star", "unit": ""},
    ])
    parsed = parse_plan(plan)
    points = parsed.generate_full_params()
    assert len(points) == 2  # 2 variable values
    for pt in points:
        assert pt["Winding_Connection"] == "Star", "String enum must be preserved"
        assert pt["Airgap"] == 1.0
        assert "Magnet_Thickness" in pt
    print("PASS: test_generate_full_params_preserves_strings")


def test_empty_fixed_params():
    """Test plan with no fixed params (boundary case)."""
    plan = _make_plan([])
    ok, errors = validate_plan_dict(plan, require_model_path=False)
    assert ok
    parsed = parse_plan(plan)
    points = parsed.generate_full_params()
    assert len(points) == 2
    print("PASS: test_empty_fixed_params")


def test_ai_generated_plan_regression():
    """Full regression test simulating the exact AI-generated plan that failed.

    The AI generator produces Winding_Connection as a string enum. This test
    ensures such plans can be parsed, validated, and expanded without the
    'could not convert string to float' error.
    """
    ai_plan = {
        "plan_id": "ai-generated-test",
        "plan_version": "2.0",
        "created_at": "2026-08-30T12:00:00",
        "topology": "SSSR",
        "model_path": "models/MARS-12S10P_SSSR.mot",
        "fixed_params": [
            {"name": "Outer_Rotor_Diameter", "value": 100.0, "unit": "mm"},
            {"name": "Inner_Rotor_Diameter", "value": 40.0, "unit": "mm"},
            {"name": "Airgap", "value": 1.0, "unit": "mm"},
            {"name": "Slot_Number", "value": 12, "unit": ""},
            {"name": "Pole_Number", "value": 10, "unit": ""},
            {"name": "Winding_Connection", "value": "Star", "unit": ""},
            {"name": "Cooling_Type", "value": "Natural Convection", "unit": ""},
            {"name": "CurrentDefinition", "value": "Peak", "unit": ""},
            {"name": "DCBusVoltage", "value": 48.0, "unit": "V"},
        ],
        "variables": [
            {"name": "Magnet_Thickness", "values": [3, 4, 5], "unit": "mm"},
            {"name": "Stator_Outer_Diameter", "values": [80, 90, 100], "unit": "mm"},
        ],
        "cases": [{"id": "default", "name": "Default", "params": {}}],
        "output_metrics": ["tavg_nm", "ripple_pct", "efficiency_pct"],
    }

    # 1. Parse should not raise
    parsed = parse_plan(ai_plan)
    assert len(parsed.fixed_params) == 9
    assert len(parsed.variables) == 2

    # 2. Validate should pass
    ok, errors = validate_plan_dict(ai_plan, require_model_path=False)
    assert ok, f"Validation failed: {errors}"

    # 3. Generate points should preserve string enums
    points = parsed.generate_full_params()
    assert len(points) == 9  # 3 * 3
    for pt in points:
        assert pt["Winding_Connection"] == "Star"
        assert pt["Cooling_Type"] == "Natural Convection"
        assert pt["CurrentDefinition"] == "Peak"
        assert isinstance(pt["Outer_Rotor_Diameter"], float)
        assert isinstance(pt["Slot_Number"], int)

    print("PASS: test_ai_generated_plan_regression")


def main():
    tests = [
        test_fixed_param_string_enum,
        test_fixed_param_numeric_values,
        test_fixed_param_to_dict_roundtrip,
        test_validate_plan_with_string_enums,
        test_generate_full_params_preserves_strings,
        test_empty_fixed_params,
        test_ai_generated_plan_regression,
    ]

    passed = 0
    failed = 0
    for test in tests:
        try:
            test()
            passed += 1
        except Exception as e:
            print(f"FAIL: {test.__name__}: {e}")
            failed += 1

    print(f"\n{'='*50}")
    print(f"Results: {passed} passed, {failed} failed, {len(tests)} total")
    if failed > 0:
        sys.exit(1)
    else:
        print("ALL TESTS PASSED")
        sys.exit(0)


if __name__ == "__main__":
    main()