evalwire/tests/test_build_subgraph.py at 5a74d655b3b03ea2999eee768aac89f6d55f3bfb · zurfjereluhmie/evalwire · GitHub

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
129
130
131
132
133
134
135
136
"""Tests for evalwire.langgraph.build_subgraph."""

from __future__ import annotations

import sys
from dataclasses import dataclass
from typing import Any

import pytest


class TestBuildSubgraph:
    """Tests for build_subgraph (requires evalwire[langgraph] extra)."""

    @pytest.fixture(autouse=True)
    def _skip_without_langgraph(self):
        pytest.importorskip("langgraph", reason="langgraph not installed")

    @pytest.fixture()
    def state_cls(self):
        @dataclass
        class State:
            value: str = ""

        return State

    def test_single_node_graph(self, state_cls: type):
        from evalwire.langgraph import build_subgraph

        def node_a(state: Any) -> dict:
            return {"value": "a"}

        compiled = build_subgraph(
            nodes=[("node_a", node_a)],
            state_cls=state_cls,
        )
        assert compiled is not None

    def test_linear_edge_wiring(self, state_cls: type):
        """START -> n1 -> n2 -> n3 -> END."""
        from evalwire.langgraph import build_subgraph

        call_order: list[str] = []

        def node_1(state: Any) -> dict:
            call_order.append("n1")
            return {"value": state.value + "1"}

        def node_2(state: Any) -> dict:
            call_order.append("n2")
            return {"value": state.value + "2"}

        def node_3(state: Any) -> dict:
            call_order.append("n3")
            return {"value": state.value + "3"}

        compiled = build_subgraph(
            nodes=[("n1", node_1), ("n2", node_2), ("n3", node_3)],
            state_cls=state_cls,
        )
        result = compiled.invoke({"value": ""})
        assert result["value"] == "123"
        assert call_order == ["n1", "n2", "n3"]

    def test_with_name_parameter(self, state_cls: type):
        from evalwire.langgraph import build_subgraph

        compiled = build_subgraph(
            nodes=[("n", lambda state: {"value": "x"})],
            state_cls=state_cls,
            name="my-graph",
        )
        assert compiled.name == "my-graph"

    def test_without_name_parameter(self, state_cls: type):
        from evalwire.langgraph import build_subgraph

        compiled = build_subgraph(
            nodes=[("n", lambda state: {"value": "x"})],
            state_cls=state_cls,
        )
        # Default name assigned by LangGraph (not None)
        assert compiled.name is not None

    def test_with_input_cls(self):
        from evalwire.langgraph import build_subgraph

        @dataclass
        class FullState:
            query: str = ""
            result: str = ""

        @dataclass
        class InputState:
            query: str = ""

        compiled = build_subgraph(
            nodes=[("n", lambda state: {"result": "done"})],
            state_cls=FullState,
            input_cls=InputState,
        )
        result = compiled.invoke({"query": "hello"})
        assert result["result"] == "done"

    def test_with_checkpointer(self, state_cls: type):
        mem_mod = pytest.importorskip(
            "langgraph.checkpoint.memory",
            reason="langgraph checkpoint not installed",
        )

        from evalwire.langgraph import build_subgraph

        compiled = build_subgraph(
            nodes=[("n", lambda state: {"value": "x"})],
            state_cls=state_cls,
            checkpointer=mem_mod.InMemorySaver(),
        )
        assert compiled is not None

    def test_empty_nodes_raises(self, state_cls: type):
        from evalwire.langgraph import build_subgraph

        with pytest.raises(IndexError):
            build_subgraph(nodes=[], state_cls=state_cls)


class TestBuildSubgraphImportError:
    def test_raises_import_error_without_langgraph(self, monkeypatch):
        """build_subgraph raises ImportError if langgraph is not installed."""
        monkeypatch.setitem(sys.modules, "langgraph", None)
        monkeypatch.setitem(sys.modules, "langgraph.graph", None)

        from evalwire.langgraph import build_subgraph

        with pytest.raises(ImportError, match="evalwire\\[langgraph\\]"):
            build_subgraph(nodes=[("n", lambda s: s)], state_cls=object)