tests/test_design_state.py

"""Tests for agents/design_state.py — state tracking and decision extraction."""

from agents.agent_flow import AgentResponse
from agents.design_state import DesignPlan, DesignState, compute_score, extract_decisions


class TestDesignState:
    def test_empty_render(self):
        state = DesignState()
        assert state.render() == ""

    def test_render_with_fields(self):
        state = DesignState(
            part_name="bracket",
            material="aluminum 6061",
            dimensions={"width": 60.0, "height": 40.0},
        )
        rendered = state.render()
        assert "bracket" in rendered
        assert "aluminum 6061" in rendered
        assert "width=60.0mm" in rendered

    def test_render_features(self):
        state = DesignState(features=["4x M6 holes", "fillet"])
        rendered = state.render()
        assert "4x M6 holes" in rendered

    def test_render_decisions_capped_at_5(self):
        state = DesignState(decisions=[f"decision {i}" for i in range(10)])
        rendered = state.render()
        assert "decision 9" in rendered
        assert "decision 4" not in rendered


class TestExtractDecisions:
    def test_extracts_material(self):
        responses = [
            AgentResponse.from_agent("engineering", "I recommend aluminum 6061 for this application.")
        ]
        state = extract_decisions(responses, DesignState())
        assert "aluminum" in state.material.lower()

    def test_extracts_dimensions_from_user(self):
        responses = []
        state = extract_decisions(responses, DesignState(), user_message="Make it 60mm wide and 40mm high")
        assert state.dimensions.get("width") == 60.0
        assert state.dimensions.get("height") == 40.0

    def test_extracts_fastener_features(self):
        responses = [
            AgentResponse.from_agent("engineering", "I'll add 4x M6 clearance holes for mounting.")
        ]
        state = extract_decisions(responses, DesignState())
        assert any("M6" in f for f in state.features)

    def test_extracts_axis_recommendation(self):
        responses = [
            AgentResponse.from_agent("cnc", "This part needs 5-axis machining due to the undercut.")
        ]
        state = extract_decisions(responses, DesignState())
        assert "5-axis" in state.axis_recommendation

    def test_extracts_part_name(self):
        responses = []
        state = extract_decisions(responses, DesignState(), user_message="I need a servo bracket with M4 holes")
        assert "servo bracket" in state.part_name.lower()

    def test_preserves_existing_state(self):
        existing = DesignState(material="steel", dimensions={"width": 50.0})
        responses = [
            AgentResponse.from_agent("engineering", "Height should be 30mm.")
        ]
        updated = extract_decisions(responses, existing, user_message="add height")
        assert updated.material == "steel"
        assert updated.dimensions.get("width") == 50.0

    def test_extracts_decisions_from_agreement(self):
        responses = [
            AgentResponse.from_agent("design", "I'd recommend an L-bracket form factor for this.")
        ]
        state = extract_decisions(responses, DesignState())
        assert len(state.decisions) > 0

    def test_no_duplicate_features(self):
        existing = DesignState(features=["4x M6 holes"])
        responses = [
            AgentResponse.from_agent("engineering", "The 4x M6 holes are properly specified.")
        ]
        updated = extract_decisions(responses, existing)
        m6_count = sum(1 for f in updated.features if "M6" in f)
        assert m6_count == 1


class TestDesignPlan:
    def test_create_from_state(self):
        state = DesignState(
            part_name="bracket",
            description="mounting bracket",
            material="aluminum 6061",
            dimensions={"width": 60.0, "height": 40.0, "depth": 20.0},
            features=["4x M6 holes"],
            constraints=["min wall 3mm"],
            axis_recommendation="3-axis",
            decisions=["use aluminum"],
        )
        plan = DesignPlan.from_state(state, confidence_score=9.0)
        assert plan.part_name == "bracket"
        assert plan.material == "aluminum 6061"
        assert plan.dimensions == {"width": 60.0, "height": 40.0, "depth": 20.0}
        assert plan.confidence_score == 9.0
        assert plan.machining_notes == []

    def test_plan_render(self):
        plan = DesignPlan(
            part_name="bracket",
            description="test",
            material="aluminum 6061",
            dimensions={"width": 60.0},
            features=["4x M6 holes"],
            constraints=[],
            axis_recommendation="3-axis",
            machining_notes=["No undercuts"],
            confidence_score=9.0,
        )
        rendered = plan.render_approved()
        assert "APPROVED DESIGN PLAN" in rendered
        assert "aluminum 6061" in rendered
        assert "No undercuts" in rendered

    def test_plan_notes_default_empty(self):
        plan = DesignPlan(part_name="test")
        assert plan.notes == ""

    def test_plan_notes_in_render(self):
        plan = DesignPlan(
            part_name="bracket",
            material="aluminum",
            notes="Check if 304 is overkill",
        )
        rendered = plan.render_approved()
        assert "User Notes" in rendered
        assert "Check if 304 is overkill" in rendered

    def test_plan_notes_empty_not_in_render(self):
        plan = DesignPlan(part_name="bracket", material="aluminum", notes="")
        rendered = plan.render_approved()
        assert "User Notes" not in rendered

    def test_plan_from_state_no_notes(self):
        state = DesignState(part_name="bracket", material="steel")
        plan = DesignPlan.from_state(state, confidence_score=5.0)
        assert plan.notes == ""


class TestComputeScore:
    def test_empty_state_scores_zero(self):
        assert compute_score(DesignState()) == 0.0

    def test_material_scores_3(self):
        state = DesignState(material="aluminum")
        assert compute_score(state) == 3.0

    def test_full_state_above_threshold(self):
        state = DesignState(
            part_name="bracket",
            description="test bracket",
            material="aluminum 6061",
            dimensions={"width": 60.0, "height": 40.0, "depth": 20.0},
            features=["4x M6 holes"],
            constraints=["min wall 3mm"],
            axis_recommendation="3-axis",
        )
        score = compute_score(state)
        assert score >= 8.0

    def test_dimension_cap_at_4(self):
        state = DesignState(dimensions={
            "width": 60, "height": 40, "depth": 20,
            "length": 100, "diameter": 10, "radius": 5,
        })
        score = compute_score(state)
        assert score == 4.0

    def test_feature_cap_at_4(self):
        state = DesignState(features=["a", "b", "c", "d", "e", "f"])
        score = compute_score(state)
        assert score == 4.0

    def test_constraint_cap_at_2(self):
        state = DesignState(constraints=["a", "b", "c", "d"])
        score = compute_score(state)
        assert score == 2.0


class TestDesignStatePhase:
    def test_default_phase_exploring(self):
        state = DesignState()
        assert state.phase == "exploring"
        assert state.plan is None

    def test_phase_serialization(self):
        plan = DesignPlan(
            part_name="b", description="", material="steel",
            dimensions={}, features=[], constraints=[],
            axis_recommendation="", machining_notes=[],
            confidence_score=5.0,
        )
        state = DesignState(phase="planning", plan=plan)
        d = state.model_dump()
        assert d["phase"] == "planning"
        assert d["plan"]["material"] == "steel"

    def test_roundtrip_from_dict(self):
        state = DesignState(phase="approved", material="brass")
        d = state.model_dump()
        restored = DesignState(**d)
        assert restored.phase == "approved"
        assert restored.material == "brass"