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

> Learn to write unit tests for Embabel agents. Extend EmbabelMockitoIntegrationTest, stub LLM calls, invoke goals, and verify results and interactions for robust testing.

- Repository: [Embabel/embabel-agent](https://github.com/embabel/embabel-agent)
- Tags: how-to-guide
- Published: 2026-08-08

---

**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`](https://github.com/embabel/embabel-agent/blob/main/EmbabelMockitoIntegrationTest.java) (located at [`embabel-agent-test-support/embabel-agent-test/src/test/java/com/embabel/agent/test/integration/EmbabelMockitoIntegrationTest.java`](https://github.com/embabel/embabel-agent/blob/main/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`](https://github.com/embabel/embabel-agent/blob/main/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`](https://github.com/embabel/embabel-agent/blob/main/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.

```java
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.

```java
@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.

```java
@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()`.

```java
    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.

```java
    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.

```java
    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`](https://github.com/embabel/embabel-agent/blob/main/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`](https://github.com/embabel/embabel-agent/blob/main/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`](https://github.com/embabel/embabel-agent/blob/main/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.