DEV Community

Cover image for Stop invoking the whole graph to test one node: unit tests and checkpoints for my LangGraph agent
Veaceslav
Veaceslav

Posted on Edited on

Stop invoking the whole graph to test one node: unit tests and checkpoints for my LangGraph agent

After part 3, my LangGraph graph worked. It also had exactly one test: run it, look at the verdict, shrug. Every edit to judge_node meant rerunning the whole graph and hoping the retries would still stop after three calls, and hope is a lousy CI stage. The spend policy of my agent, the code that decides when the copilot is allowed to pay for LLM calls, got audited the way I check soup: with a fork and a squint.

A comment on part 3 named the two missing pieces: stop testing through the framework, because the nodes are pure functions and deserve unit tests, and add checkpoints between runs. This part is both. What I have now: thirteen tests that run in 0.26 seconds without a single LLM call, and a checkpoint that parks the graph right before the model starts costing money.

The JobCopilot mirror is unreachable from this machine, so the code below is a checked template rebuilt from the published contracts of parts 1-3 (AgentState, route_after_parse, judge_node(state, chat), build_graph(chat)). Before publishing, align node names, prompt wording, and the THRESHOLD = 60 / MAX_CALLS = 3 constants with your repo. The numbers below come from a real run in my sandbox (langgraph 1.2.20, pydantic 2.14.3, Python 3.14).*

The router is twelve lines, so it gets a truth table

The router from part 3, route_after_parse, decides between ok, retry, and give_up. It is the entire economy of my agent in one pure function: no LLM, no framework, no network. A truth-table test costs nothing to write:

# test_copilot.py
import pytest

from nodes import route_after_parse
from state import MAX_CALLS, AgentState


@pytest.mark.parametrize(
    "error, attempt, expected",
    [
        (None, 0, "ok"),
        ("schema: 1 validation error", 1, "retry"),
        ("schema: 1 validation error", 2, "retry"),
        ("schema: 1 validation error", MAX_CALLS, "give_up"),
    ],
)
def test_router_branches(error, attempt, expected):
    state = AgentState(error=error, attempt=attempt)
    assert route_after_parse(state) == expected


def test_router_budget_boundary():
    """The retry budget runs out exactly at attempt == MAX_CALLS, not one past it."""
    assert route_after_parse(AgentState(error="x", attempt=MAX_CALLS - 1)) == "retry"
    assert route_after_parse(AgentState(error="x", attempt=MAX_CALLS)) == "give_up"
Enter fullscreen mode Exit fullscreen mode

That boundary test is the whole point of having MAX_CALLS as a constant instead of a magic number inside the loop. Off-by-one errors in retry budgets are invisible in a happy-path demo: the model cooperates, the router never reaches the edge. This test visits the edge on purpose, so the day I change 3 to 4, it fails loudly instead of quietly burning one extra paid call per job description.

Coming from React, this was my first genuinely comfortable pytest moment. I test pure functions the same way I test reducers.

Node tests: the fake chat, and what the prompts actually say

judge_node(state, chat) takes its LLM callable as an argument, a design decision from part 3 that pays off here. Unit testing it is just handing it a fake, and my fake records the prompt it was sent:

def recording_chat(raw: str):
    sent = {}

    def chat(prompt: str) -> str:
        sent["prompt"] = prompt
        return raw

    return chat, sent
Enter fullscreen mode Exit fullscreen mode

The LLM call becomes free, deterministic, and inspectable. A test can ask what the node actually asked the model, not only what it did with the answer:

def test_first_pass_asks_about_the_jd():
    chat, sent = recording_chat('{"match_score": 72, "missing_skills": ["Terraform"]}')
    state = AgentState(jd_text="Senior Frontend (React, TS)", context="<chunks>")

    patch = judge_node(state, chat)

    assert "Senior Frontend (React, TS)" in sent["prompt"]
    assert "RESUME EXCERPT" in sent["prompt"]
    assert "PREVIOUS" not in sent["prompt"]  # no repair talk on a fresh judge


def test_retry_prompt_quotes_the_model_back():
    chat, sent = recording_chat('{"match_score": 70, "missing_skills": []}')
    state = AgentState(
        error="Expecting value: line 1 column 1 (char 0)",
        last_raw="Sure! The match is quite good, I'd say around 70 out of 100.",
    )

    judge_node(state, chat)

    assert "PARSER ERROR" in sent["prompt"]
    assert "around 70 out of 100" in sent["prompt"]  # its prose, quoted back
Enter fullscreen mode Exit fullscreen mode

The second test is the repair trick from part 1 turned into a regression check. If a future refactor "simplifies" the repair prompt and stops quoting the model's own prose back at it, nothing will crash: the graph reruns, the verdict looks plausible, and repair quality rots in silence. The test can't grade repair quality, but it can check the mechanism survived.

The test that actually caught a design mistake of mine is about patch shape:

def test_judge_returns_a_state_patch_nothing_else():
    chat, _ = recording_chat('{"match_score": 72, "missing_skills": ["Terraform"]}')
    patch = judge_node(AgentState(jd_text="x", context="y"), chat)

    assert patch["attempt"] == 1
    assert patch["match_score"] == 72
    assert patch["verdict"] == "match"
    assert patch["error"] is None
    assert set(patch) == {"last_raw", "attempt", "match_score", "missing_skills", "verdict", "error"}
Enter fullscreen mode Exit fullscreen mode

My first draft of that last line asserted that the error case writes verdict: None into the patch. It failed, because the error patch carries no verdict key at all. Writing the assertion forced me to confront something the design already knew: an undecided verdict should be absent, not None. The router reads state.error, so a verdict: None in the patch is a lie the next reader has to interpret. The test made the "no decision, no key" rule explicit, which is the one thing a comment cannot do.

Match this anecdote to how judge_node actually evolved in your repo; the assertion above ("verdict" not in patch on a failed parse) follows the part 3 code published earlier. Also recheck the ParseOk field name: part 1 used ParseOk(match=JobMatch(...)), the part 3 draft read result.data.*

Testing the loop through the graph, still no LLM

Node tests check the parts. One test per promise that belongs to the graph itself: the repair loop recovers, and the retry budget stops.

import json

from graph import build_graph
from state import AgentState, MAX_CALLS


def test_budget_loop_ends_cleanly():
    calls = {"n": 0}

    def chat(prompt: str) -> str:
        calls["n"] += 1
        return "I'd say around 70 out of 100."  # prose, never JSON

    final = build_graph(chat).invoke(AgentState(jd_text="Junior Python Backend"))

    assert final["verdict"] == "error"  # invoke returns state as a dict
    assert final["attempt"] == MAX_CALLS
    assert calls["n"] == MAX_CALLS  # the loop stopped itself


def test_repair_loop_recovers():
    calls = {"n": 0}

    def chat(prompt: str) -> str:
        calls["n"] += 1
        if calls["n"] == 1:
            return "Sure! The match is quite good, I'd say around 70 out of 100."
        return json.dumps({"match_score": 70, "missing_skills": ["Terraform"]})

    final = build_graph(chat).invoke(AgentState(jd_text="Senior Frontend Engineer (React, TS)"))

    assert final["verdict"] == "match"  # invoke returns state as a dict
    assert final["match_score"] == 70
    assert final["attempt"] == 2
Enter fullscreen mode Exit fullscreen mode

Both are end-to-end runs and both are free: a fake chat goes in, assertions come out. The division of labor I settled on: node tests own the details (which prompt fired, the exact shape of the patch), graph tests own the promises between nodes (the budget exhausts, the loop recovers, fail exits without raising).

The whole suite: thirteen tests, 0.26 seconds. One real LLM round trip takes longer, so the suite can run before every commit, which a suite with an LLM inside never could.

Two honest notes from the way there. First, final.verdict raises AttributeError: app.invoke() returns the merged state values as a plain dict, not my pydantic AgentState. The framework hands you channel values; the attribute access I reached for was my model, not the graph's interface. Second, invoke(AgentState(...)) is input, not a stored run: every invoke is a fresh run unless a thread_id says otherwise. Which leads straight into checkpointing.

The dict-return behavior is checked against langgraph 1.2.20; reconfirm on your pinned version.*

Checkpointing: state leaves the stack

With a checkpointer, LangGraph saves state after every node. A run stops being a black box that either finishes or throws. You can stand between nodes and look around:

import json

from langgraph.checkpoint.memory import MemorySaver

from graph import build_graph
from state import AgentState


def compliant_chat(score: int = 70):
    def chat(prompt: str) -> str:
        return json.dumps({"match_score": score, "missing_skills": ["Terraform"]})

    return chat


JD = "Senior Frontend Engineer (React, TypeScript, 5+ years)"

app = build_graph(compliant_chat(), checkpointer=MemorySaver(), interrupt_before=["judge"])
config = {"configurable": {"thread_id": "jd-42"}}
Enter fullscreen mode Exit fullscreen mode

Two settings changed from part 3's build_graph: checkpointer=MemorySaver() turns on state snapshots (in-memory is fine for a batch run inside one process), and interrupt_before=["judge"] tells the graph to stop right at the LLM node's door. Run it, inspect, resume:

1. run stops before the LLM call
2. next: ('judge',) | values: {'jd_text': 'Senior Frontend Engineer (React, TypeScript, 5+ years)',
   'context': '<top-8 resume chunks>', 'attempt': 0, 'missing_skills': []}
3. resume finished: next: () | values: {..., 'attempt': 1,
   'match_score': 70, 'missing_skills': ['Terraform'], 'verdict': 'match'}
Enter fullscreen mode Exit fullscreen mode

Step 2 is the part I did not expect to enjoy this much. next: ('judge',) is the graph parked, and it tells me where it is:

stopped = app.invoke(AgentState(jd_text=JD), config)
snap = app.get_state(config)
print(snap.next)    # ('judge',) parked at the LLM door
print(snap.values)  # jd + resume chunks, attempt: 0, no score yet
Enter fullscreen mode Exit fullscreen mode

The snapshot holds the job description and the extracted resume chunks with zero LLM calls spent. In copilot terms: I can check exactly what the judge will see before the judge sees it. That is a human-in-the-loop door I never planned to build; it fell out of the checkpoint for free.

Then app.invoke(None, config), where None means "continue, no new input", and step 3 confirms the resume worked: attempt: 1, not 2. The retrieve node did not run again. The run picked up exactly where it parked, including the LLM call it owed. And the part that makes this actual infrastructure: threads are isolated.

fresh = app.invoke(AgentState(jd_text=JD), {"configurable": {"thread_id": "jd-43"}})
snap = app.get_state(config)
print(fresh)        # jd-43's own run, fresh attempt
print(snap.values)  # jd-42's checkpoint, exactly where we left it
Enter fullscreen mode Exit fullscreen mode

thread_id is the boundary between "this job description's run" and "every other run". When JobCopilot chews through a night batch of job descriptions, each gets its own thread_id, and every run stays inspectable after the fact, including the ones that errored. It also opens the door to durable storage later: swap MemorySaver for a persistent checkpointer and a crashed batch run resumes instead of restarting from zero. That one is on my todo list, not in my code, and this post only claims runs I actually executed.

One more honest surprise: on step 2, snap.values["match_score"] raised KeyError. Not every state field lands in the snapshot. With my pydantic state the checkpoint carries only fields written at least once, and untouched None fields are simply not there.

The outputs above are verbatim from my test run (MemorySaver + interrupt_before + pydantic state, langgraph 1.2.20). Worth repeating on your pinned langgraph.*

What changed this week

My agent's spend policy went from "run it and look" to a test suite I trust. The one-line version: test the router as a truth table, test the nodes with a fake chat, test the budget through the graph, and let checkpoints turn runs into resumable threads.

Three things this React dev learning Python took from it:

  • The patch-shape assertion turned out to be the best test in the file. It pins a contract: an undecided verdict is absent, not None. Comments describe intent; assertions enforce it.
  • The fake chat makes the LLM deterministic. Prompt content, the repair mechanism, the retry budget: all of it tested in a quarter of a second.
  • interrupt_before plus thread_id turned the graph into something more than a function call: a stored, inspectable process. next: ('judge',) might be my favorite bit of Python output so far.

Next up: the parts of the copilot that talk to the real world, or swapping the in-memory checkpointer for persistent storage if something breaks first. In case you missed it, part 3: Stop calling functions: my first LangGraph graph, where state does the calling.

Follow along if you want to watch the rest of the copilot get built, including the parts that break. Build log plus daily code at t.me/novamind_hub.

How do you test nodes that wrap an LLM call: fakes, or recorded real responses? And has checkpointing paid off in your graph?

Top comments (0)