File size: 4,358 Bytes
ba5110e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
"""
Test cases for LangGraph agent workflow.
Tests state, graph compilation, and routing logic.
"""
import pytest
from backend.agent.state import AgentState
from backend.agent.graph import build_graph, agent_graph
from backend.agent.nodes import should_use_tool


class TestAgentState:
    """Test suite for agent state definitions."""

    def test_state_structure(self):
        """TC-LG-001: AgentState should have all required fields."""
        state: AgentState = {
            "messages": [],
            "session_id": "test-session",
            "current_model": "openai/gpt-oss-120b",
            "tool_retry_count": 0,
            "code_correction_count": 0,
            "wolfram_retry_count": 0,
            "error_message": None,
            "should_fallback": False,
            "image_data": None,
        }
        assert state["session_id"] == "test-session"
        assert state["current_model"] == "openai/gpt-oss-120b"

    def test_state_model_options(self):
        """TC-LG-002: Model should be one of the allowed values."""
        valid_models = ["openai/gpt-oss-120b", "openai/gpt-oss-20b"]
        state: AgentState = {
            "messages": [],
            "session_id": "test",
            "current_model": "openai/gpt-oss-120b",
            "tool_retry_count": 0,
            "code_correction_count": 0,
            "wolfram_retry_count": 0,
            "error_message": None,
            "should_fallback": False,
            "image_data": None,
        }
        assert state["current_model"] in valid_models


class TestGraphCompilation:
    """Test suite for LangGraph compilation."""

    def test_graph_compiles(self):
        """TC-LG-003: Graph should compile without errors."""
        graph = build_graph()
        assert graph is not None

    def test_agent_graph_exists(self):
        """TC-LG-004: Pre-compiled agent_graph should exist."""
        assert agent_graph is not None


class TestRoutingLogic:
    """Test suite for graph routing decisions."""

    def test_route_to_fallback_when_should_fallback(self):
        """TC-LG-005: Should route to fallback when flag is set."""
        state: AgentState = {
            "messages": [],
            "session_id": "test",
            "current_model": "openai/gpt-oss-120b",
            "tool_retry_count": 0,
            "code_correction_count": 0,
            "wolfram_retry_count": 0,
            "error_message": "Test error",
            "should_fallback": True,
            "image_data": None,
        }
        result = should_use_tool(state)
        assert result == "fallback"

    def test_route_to_tool_when_pending(self):
        """TC-LG-006: Should route to tool when pending tool exists."""
        state: AgentState = {
            "messages": [],
            "session_id": "test",
            "current_model": "openai/gpt-oss-120b",
            "tool_retry_count": 0,
            "code_correction_count": 0,
            "wolfram_retry_count": 0,
            "error_message": None,
            "should_fallback": False,
            "image_data": None,
            "_pending_tool": "wolfram",
        }
        result = should_use_tool(state)
        assert result == "tool"

    def test_route_to_format_when_tool_result(self):
        """TC-LG-007: Should route to format when tool result exists."""
        state: AgentState = {
            "messages": [],
            "session_id": "test",
            "current_model": "openai/gpt-oss-120b",
            "tool_retry_count": 0,
            "code_correction_count": 0,
            "wolfram_retry_count": 0,
            "error_message": None,
            "should_fallback": False,
            "image_data": None,
            "_tool_result": "x = 5",
        }
        result = should_use_tool(state)
        assert result == "format"

    def test_route_to_end_when_complete(self):
        """TC-LG-008: Should route to end when no pending actions."""
        state: AgentState = {
            "messages": [],
            "session_id": "test",
            "current_model": "openai/gpt-oss-120b",
            "tool_retry_count": 0,
            "code_correction_count": 0,
            "wolfram_retry_count": 0,
            "error_message": None,
            "should_fallback": False,
            "image_data": None,
        }
        result = should_use_tool(state)
        assert result == "end"