在您完成了 LangGraph Agent 的原型设计后,下一步自然是为它添加测试。本指南介绍了一些在编写单元测试时可以使用的有用模式。 请注意,本指南是 LangGraph 特定的,涵盖了具有自定义结构的图的场景——如果您刚刚开始,请查看使用 LangChain 内置 create_agent测试

前置条件

首先,确保您已安装 pytest
$ pip install -U pytest

开始使用

由于许多 LangGraph Agent 依赖于状态,一个有用的模式是在每个使用它的测试之前创建您的图,然后在测试中使用新的检查点保存器实例编译它。 下面的示例展示了如何通过一个简单的线性图来实现这一点,该图经过 node1node2。每个节点更新单个状态键 my_key
import pytest

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

def create_graph() -> StateGraph:
    class MyState(TypedDict):
        my_key: str

    graph = StateGraph(MyState)
    graph.add_node("node1", lambda state: {"my_key": "hello from node1"})
    graph.add_node("node2", lambda state: {"my_key": "hello from node2"})
    graph.add_edge(START, "node1")
    graph.add_edge("node1", "node2")
    graph.add_edge("node2", END)
    return graph

def test_basic_agent_execution() -> None:
    checkpointer = MemorySaver()
    graph = create_graph()
    compiled_graph = graph.compile(checkpointer=checkpointer)
    result = compiled_graph.invoke(
        {"my_key": "initial_value"},
        config={"configurable": {"thread_id": "1"}}
    )
    assert result["my_key"] == "hello from node2"

测试单个节点和边

编译后的 LangGraph Agent 将每个单独节点的引用公开为 graph.nodes。您可以利用这一点来测试 Agent 中的单个节点。请注意,这将绕过编译图时传入的任何检查点保存器:
import pytest

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

def create_graph() -> StateGraph:
    class MyState(TypedDict):
        my_key: str

    graph = StateGraph(MyState)
    graph.add_node("node1", lambda state: {"my_key": "hello from node1"})
    graph.add_node("node2", lambda state: {"my_key": "hello from node2"})
    graph.add_edge(START, "node1")
    graph.add_edge("node1", "node2")
    graph.add_edge("node2", END)
    return graph

def test_individual_node_execution() -> None:
    # 在本例中将被忽略
    checkpointer = MemorySaver()
    graph = create_graph()
    compiled_graph = graph.compile(checkpointer=checkpointer)
    # 仅调用节点 1
    result = compiled_graph.nodes["node1"].invoke(
        {"my_key": "initial_value"},
    )
    assert result["my_key"] == "hello from node1"

部分执行

对于由较大图组成的 Agent,您可能希望测试 Agent 内部的执行路径,而不是端到端的整个流程。在某些情况下,从语义上讲,将这些部分重构为子图 可能更有意义,您可以像普通图一样单独调用它们。 但是,如果您不希望更改 Agent 图的整体结构,您可以使用 LangGraph 的持久化机制来模拟 Agent 在所需部分开始之前暂停的状态,并将在所需部分结束时再次暂停。步骤如下:
  1. 使用检查点保存器编译您的 Agent(内存中的检查点保存器 InMemorySaver 对于测试来说已经足够)。
  2. 使用 as_node 参数调用 Agent 的 update_state 方法,该参数设置为您想要开始测试的节点之前的节点名称。
  3. 使用相同的 thread_id 调用您的 Agent,并设置 interrupt_after 参数为您想要停止的节点名称。
以下是仅执行线性图中第二个和第三个节点的示例:
import pytest

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver

def create_graph() -> StateGraph:
    class MyState(TypedDict):
        my_key: str

    graph = StateGraph(MyState)
    graph.add_node("node1", lambda state: {"my_key": "hello from node1"})
    graph.add_node("node2", lambda state: {"my_key": "hello from node2"})
    graph.add_node("node3", lambda state: {"my_key": "hello from node3"})
    graph.add_node("node4", lambda state: {"my_key": "hello from node4"})
    graph.add_edge(START, "node1")
    graph.add_edge("node1", "node2")
    graph.add_edge("node2", "node3")
    graph.add_edge("node3", "node4")
    graph.add_edge("node4", END)
    return graph

def test_partial_execution_from_node2_to_node3() -> None:
    checkpointer = MemorySaver()
    graph = create_graph()
    compiled_graph = graph.compile(checkpointer=checkpointer)
    compiled_graph.update_state(
        config={
          "configurable": {
            "thread_id": "1"
          }
        },
        # 传入节点 2 的状态 - 模拟节点 1 结束时的状态
        values={"my_key": "initial_value"},
        # 将保存的状态更新为好像来自节点 1
        # 执行将在节点 2 恢复
        as_node="node1",
    )
    result = compiled_graph.invoke(
        # 通过传递 None 恢复执行
        None,
        config={"configurable": {"thread_id": "1"}},
        # 在节点 3 之后停止,以便节点 4 不运行
        interrupt_after="node3",
    )
    assert result["my_key"] == "hello from node3"