How to Write Unit Tests for Embabel Agents: A Complete Guide

Unit testing Embabel agents requires extending EmbabelMockitoIntegrationTest to boot an in-memory AgentPlatform, stubbing LLM calls via Mockito, invoking goals through AgentInvocation, and verifying both results and interactions.

Embabel agents run on the AgentPlatform runtime, where annotated methods (@Action, @AchievesGoal) define behavior and the AgentInvocation API triggers execution. In the embabel/embabel-agent repository, the testing framework isolates large language model (LLM) dependencies behind the Ai façade, allowing you to write fast, deterministic unit tests that never hit real model endpoints.

Core Architecture of Embabel Agent Testing

Understanding the three pillars of the test framework helps you write robust assertions against agent orchestration logic.

AgentPlatform and the Test Lifecycle

The AgentPlatform acts as a lightweight container that holds the agent registry and manages AgentProcess lifecycles. In EmbabelMockitoIntegrationTest.java (located at embabel-agent-test-support/embabel-agent-test/src/test/java/com/embabel/agent/test/integration/EmbabelMockitoIntegrationTest.java), the base class automatically starts an in-memory platform before each test and registers a Mockito-based mock for the LLM provider. This setup eliminates external network calls while preserving the full agent execution path.

AgentInvocation and Goal Resolution

AgentInvocation, defined in embabel-agent-api/src/main/kotlin/com/embabel/agent/api/invocation/AgentInvocation.kt, provides a static factory method create() that builds a typed invocation for a specific goal-output class. When you call invoke(), the platform resolves the correct agent, wires the Ai instance, and executes the method annotated with @AchievesGoal. This abstraction lets you test the complete workflow from input to exported result without manual dependency injection.

The Ai Façade for LLM Stubbing

The Ai helper (found in embabel-agent-api/src/main/kotlin/com/embabel/agent/api/common/Ai.kt) decorates LLM calls with configuration options like withLlm(), withAutoLlm(), and withPromptContributor(). In tests, the underlying LLM implementation is a Mockito mock, enabling you to intercept calls to createObject and generateText using the stubbing methods provided by EmbabelMockitoIntegrationTest.

Setting Up the Test Harness

Every unit test for an Embabel agent follows a consistent initialization pattern. Extend the abstract base class to inherit platform management and stubbing utilities.

import com.embabel.agent.test.integration.EmbabelMockitoIntegrationTest;

class MyAgentTest extends EmbabelMockitoIntegrationTest {
    // Test methods inherit access to agentPlatform and Mockito stubs
}

This base class provides:

• Automatic platform boot: An in-memory AgentPlatform instance available as agentPlatform
• Mock LLM registration: Pre-configured stubs accessible via whenCreateObject() and whenGenerateText()
• Verification helpers: Methods like verifyCreateObjectMatching() and verifyNoMoreInteractions() for assertion

How to Write Unit Tests for Embabel Agents

A complete test flow consists of five distinct steps, from input creation to interaction verification.

1. Define the Agent Under Test

Create an inner class or separate file with the @Agent annotation, action methods, and goal-export definitions. Use Java records for DTOs to ensure immutability and clear contracts.

@Agent(description = "Blocking agent for testing")
private static class BlockingTestAgent {

    public record Story(String text) {}
    public record ReviewedStory(Story story, String review) {}

    @Action
    Story craftStory(UserInput userInput, Ai ai) {
        return ai.withLlm(LlmOptions.withDefaultLlm().withTemperature(0.9))
                 .createObject("Craft a short story about: " + userInput.getContent(), 
                               Story.class);
    }

    @AchievesGoal(
        description = "Story has been crafted and reviewed",
        export = @Export(remote = true, name = "writeAndReviewStory"))
    @Action
    ReviewedStory reviewStory(Story story, Ai ai) {
        String review = ai.withAutoLlm()
                         .generateText("Review this story: " + story.text());
        return new ReviewedStory(story, review);
    }
}

2. Stub LLM Responses

Use the inherited stubbing methods to define deterministic responses for specific prompt patterns. This prevents non-deterministic AI behavior from breaking assertions.

@Test
void shouldExecuteCompleteWorkflow() {
    var input = new UserInput("Write about artificial intelligence");

    // Stub object generation
    whenCreateObject(s -> s.contains("Craft a short story"), 
                     BlockingTestAgent.Story.class)
        .thenReturn(new BlockingTestAgent.Story("AI transforms our world..."));

    // Stub text generation
    whenGenerateText(s -> s.contains("Review this story"))
        .thenReturn("Excellent exploration of AI themes.");

3. Invoke the Agent Goal

Construct an AgentInvocation targeting the goal's return type, then pass the UserInput to invoke().

    var invocation = AgentInvocation.create(agentPlatform, 
                                            BlockingTestAgent.ReviewedStory.class);
    var result = invocation.invoke(input);

4. Assert the Outcome

Verify that the returned DTO contains expected data structures and business logic results.

    assertNotNull(result);
    assertTrue(result.story().text().contains("AI transforms"));
    assertEquals("Excellent exploration of AI themes.", result.review());

5. Verify LLM Interactions

Confirm that the agent used the expected prompts and configuration options, such as temperature settings or tool groups.

    verifyCreateObjectMatching(
        prompt -> prompt.contains("Craft a short story"),
        BlockingTestAgent.Story.class,
        llm -> llm.getLlm().getTemperature() != null && 
               llm.getLlm().getTemperature() == 0.9 &&
               llm.getToolGroups().isEmpty());
    
    verifyGenerateTextMatching(
        prompt -> prompt.contains("Review this story"));
    
    verifyNoMoreInteractions();
}

Verifying LLM Configuration and Prompts

Beyond checking return values, Embabel tests should validate that the agent configured the LLM correctly. The verifyCreateObjectMatching and verifyGenerateTextMatching methods accept three arguments:

1. A predicate to match the prompt string
2. The expected return class (for typed object creation)
3. A consumer to inspect LlmOptions (found in embabel-agent-api/src/main/kotlin/com/embabel/common/ai/model/LlmOptions.kt)

This pattern ensures that production code calling withLlm(LlmOptions.withDefaultLlm().withTemperature(0.9)) actually passes those options to the underlying provider, preventing regression in LLM parameter tuning.

Summary

• Extend EmbabelMockitoIntegrationTest to inherit an in-memory AgentPlatform and Mockito-based LLM stubs.
• Stub deterministically using whenCreateObject() and whenGenerateText() to avoid non-deterministic AI responses in CI/CD pipelines.
• Invoke via AgentInvocation.create() to test the full orchestration flow from @Action methods to @AchievesGoal exports.
• Verify interactions with verifyCreateObjectMatching() to assert prompt content and LlmOptions configuration like temperature and max tokens.
• Use Java records for DTOs to maintain type safety between agent actions and test assertions.

Frequently Asked Questions

How do I test streaming agents that return partial results?

Use the EmbabelMockitoIntegrationTestStreamingTest class (located at embabel-agent-test-support/embabel-agent-test/src/test/java/com/embabel/agent/test/integration/EmbabelMockitoIntegrationTestStreamingTest.java) as a reference. The pattern remains identical—stub whenGenerateText() with multiple responses or use thenAnswer() to simulate incremental chunks—but you invoke the goal using the same AgentInvocation API and assert on the accumulated result.

Can I verify specific LLM parameters like temperature and model name?

Yes. Pass a lambda to verifyCreateObjectMatching() or verifyGenerateTextMatching() that inspects the LlmOptions object. For example, verify that llm.getLlm().getTemperature() == 0.9 or check llm.getToolGroups() to ensure no tools were passed, as shown in the EmbabelMockitoIntegrationTestBlockingTest.java example.

Do I need to mock the entire AgentPlatform manually?

No. The EmbabelMockitoIntegrationTest base class handles platform initialization and mock registration automatically. Your test only needs to define the agent class and stub specific LLM interactions; the framework manages lifecycle cleanup and verification state between tests.

What is the difference between createObject and generateText in tests?

createObject is used when the agent calls ai.createObject() to generate a structured Java object (like a record), while generateText is used for raw string output from ai.generateText(). Always match your stubbing method to the actual API call in the agent action, and use the corresponding verification method to assert the interaction occurred with the correct prompt pattern.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}