"""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"