Stub the Model, Assert the Plan: Unit-Testing Embabel Agents

The thing that makes agent development sane isn't a clever prompt. It's a test suite that lets you refactor without clicking through the app and burning API credits on every change. In Embabel, the JVM agent framework, the non-deterministic part — the LLM call — is the only part you can't assert on directly. Stub that out, and everything else (routing, plan wiring, the prompt you actually built) is ordinary deterministic code. This is how I test it.

Two layers, two tools

Embabel agents are Spring beans made of typed @Action methods, wired into a plan by a planner. That gives you two natural test layers, and Embabel ships a helper for each:

• Unit — call one @Action as a plain method with a fake context. No Spring, no planner. Milliseconds.
• Integration — boot Spring, load the real AgentPlatform, run the whole plan through AgentInvocation. Seconds — it literally starts your app.

The economics decide the mix. In my project, five unit tests ran in 23 ms and made zero LLM calls. Each integration test took ~130 seconds, because it stands up the full context. So: push logic and routing down to unit tests, keep a thin layer of integration tests for the wiring.

Unit tests: FakeOperationContext

Every action takes an Ai handle that Embabel injects. In a test you build a fake one, tell it what the model should "say", then call the method directly. Here's a branching agent that classifies a support ticket and routes it into one of three @State handlers. I want to assert the routing without ever calling a model:

private SupportTriageAgent.Triaged triageWith(String llmVerdict, String ticket) {
    var context = FakeOperationContext.create();
    context.expectResponse(llmVerdict);   // the classifier's generateText(...) returns this
    return new SupportTriageAgent()
            .triage(new UserInput(ticket, Instant.now()), context.ai());
}

@Test
void routesToBugWhenClassifiedBug() {
    assertInstanceOf(SupportTriageAgent.Bug.class,
            triageWith("BUG", "The export button throws a 500"));
}

@Test
void fallsBackToGeneralOnUnknownVerdict() {
    // an off-script LLM answer should not crash -- it defaults to General
    assertInstanceOf(SupportTriageAgent.General.class,
            triageWith("banana", "Anything at all"));
}

expectResponse(...) queues the model's answer; context.ai() hands the action a fake that returns it. The test asserts which branch the agent took. That's the highest-value thing to lock down in a branching agent, and it costs nothing to run. Note the last test: I stub a nonsense verdict to prove the fallback path is defensive — the kind of edge case you'd never want to probe against a live model.

Assert the prompt, not just the result

The fake context also records every invocation, so you can assert on the prompt your action actually built — useful when the prompt is assembled from several inputs:

var context = FakeOperationContext.create();
var promptRunner = (FakePromptRunner) context.promptRunner();
context.expectResponse("CRITICAL");

new SupportTriageAgent()
        .triage(new UserInput("The database is on fire", Instant.now()), context.ai());

var prompt = promptRunner.getLlmInvocations().getFirst().getMessages().getFirst().getContent();
assertTrue(prompt.contains("database is on fire"), "Prompt should include the ticket text");
assertTrue(prompt.contains("CRITICAL"), "Prompt should list the CRITICAL category");

This catches a whole class of silent regressions: someone tweaks a prompt template and drops the user's input, or forgets to interpolate a field. The model would happily produce plausible garbage; the test fails loudly instead.

Integration tests: run the real plan

Unit tests never exercise the planner — the thing that decides order. For that you extend EmbabelMockitoIntegrationTest, which boots Spring and gives you a real AgentPlatform with your agents loaded. You stub the LLM with prompt-matching rules and then invoke by result type:

whenGenerateText(prompt -> prompt.contains("Classify this support ticket"))
        .thenReturn("BUG");

var invocation = AgentInvocation.create(agentPlatform, SupportTriageAgent.ResolvedTicket.class);
var resolved = invocation.invoke(new UserInput("The export button throws a 500"));

assertEquals("BUG", resolved.category());
assertEquals("ENGINEERING_TEAM", resolved.handledBy());

verifyGenerateTextMatching(prompt -> prompt.contains("Classify this support ticket"));
verifyNoMoreInteractions();

The last two lines are the point. verifyNoMoreInteractions() asserts the plan made exactly the LLM calls I expected and no others. Here it proves the agent classified the ticket, entered one state, ran that state's pure-code handler, and stopped — the branch really was exclusive, not a plan that wandered into extra calls. You're asserting the planner's behaviour, not just the final object.

Why the slow tests earn their keep

Earlier in the same session I'd refactored an agent: moved a goal from one action to another so a linear pipeline had a single terminal goal. All my unit tests stayed green — they call methods directly and don't care about goals. But one integration test went red:

java.lang.IllegalStateException: No agent with outputClass
    class ...WriteAndReviewAgent$ReviewedStory found.

That's the wiring talking. AgentInvocation selects an agent by the goal's output type, and my refactor had removed the goal that produced ReviewedStory. The unit tests couldn't have caught this, because the breakage lives in the plan graph, not in any single method. The fix was to point the test at the new terminal type and stub the extra step. Exactly the class of bug you want a machine to catch, not a user.

The takeaway

The workflow that makes this pleasant is simple: stub the model, assert the plan, refactor fearlessly. Unit tests with FakeOperationContext pin down routing and prompt construction in milliseconds. A thin layer of EmbabelMockitoIntegrationTest cases pins down the wiring the planner derives from your types. Between them, the only thing left unasserted is the model's actual words — which is the one thing you shouldn't be asserting in a unit test anyway. Box out the non-determinism, and an "AI agent" becomes a normal, testable Spring application. That's the whole trick.