# How to Write Effective Unit and Integration Tests for Custom Agents in Codebuff

> Learn to write effective unit and integration tests for custom Codebuff agents. Isolate agent logic with unit tests and verify end-to-end behavior with integration tests for robust applications.

- Repository: [Codebuff/codebuff](https://github.com/CodebuffAI/codebuff)
- Tags: best-practices
- Published: 2026-03-09

---

**Unit tests validate your agent's static `AgentDefinition` and generator logic in isolation, while integration tests execute full lifecycles via the SDK to verify tool interactions and end-to-end behavior.**

Codebuff agents are built around a typed `AgentDefinition` object and execute behavior through a generator-style `handleSteps` function. Testing these custom agents requires a two-layer approach that covers both static configuration validation and dynamic runtime interaction. This guide demonstrates how to write comprehensive tests using patterns from the official `CodebuffAI/codebuff` test suite.

## Understanding the Two-Layer Testing Strategy

Effective testing for Codebuff agents separates concerns into unit and integration layers. This division allows you to catch configuration errors quickly while ensuring complex multi-step workflows function correctly against the actual runtime engine.

### Unit Tests: Validating Agent Definitions and Logic

Unit tests focus on the static `AgentDefinition` structure and the step-generation logic within `handleSteps`. These tests run in milliseconds and verify that required fields like `id`, `model`, `toolNames`, and prompts are correctly configured without launching the full runtime.

### Integration Tests: Testing the Full Agent Lifecycle

Integration tests spawn agents via the SDK's `run` method and execute complete lifecycles against the runtime engine. These tests verify that tools like `read_files`, `write_file`, and `set_output` interact correctly, and that multi-step agents such as the *editor* or *thinker* complete their workflows successfully.

## Unit Testing Agent Definitions

The `AgentDefinition` interface in [`agents/types/agent-definition.ts`](https://github.com/CodebuffAI/codebuff/blob/main/agents/types/agent-definition.ts) declares the contract that every agent must satisfy. Unit tests should import the agent constant and assert compliance with this contract.

### Validating Required Fields and Structure

The test suite in [`agents/__tests__/editor.test.ts`](https://github.com/CodebuffAI/codebuff/blob/main/agents/__tests__/editor.test.ts) demonstrates the canonical pattern for definition testing:

```typescript
import editor from '../editor/editor'

describe('editor agent', () => {
  test('has correct id', () => {
    expect(editor.id).toBe('editor')
  })
  
  test('uses opus model by default', () => {
    expect(editor.model).toBe('anthropic/claude-opus-4.6')
  })
  
  test('has required tool names', () => {
    expect(editor.toolNames).toContain('read_files')
    expect(editor.toolNames).toContain('write_file')
  })
})

```

This approach validates that the agent exports the correct `id`, maps to the expected model, and includes necessary tools without executing any runtime code.

### Testing Model Mappings with createCodeEditor

Codebuff provides the `createCodeEditor` factory to map shorthand model names to full OpenRouter identifiers. Unit tests should verify that shortcuts like `'opus'`, `'gpt-5'`, and `'minimax'` resolve to the correct full model strings:

```typescript
import { createCodeEditor } from '../editor/factory'

describe('createCodeEditor factory', () => {
  test('maps opus shortcut to anthropic/claude-opus-4.6', () => {
    const agent = createCodeEditor('opus')
    expect(agent.model).toBe('anthropic/claude-opus-4.6')
  })
  
  test('maps gpt-5 shortcut to correct OpenRouter identifier', () => {
    const agent = createCodeEditor('gpt-5')
    expect(agent.model).toBe('openai/gpt-5')
  })
})

```

These tests ensure that model configuration remains consistent when users specify shortcuts in their agent definitions.

## Testing Agent Logic and Step Sequences

Beyond static definitions, agents implement behavior through the `handleSteps` generator function. Unit tests can invoke this generator directly to verify the sequence of yielded values without connecting to the LLM.

### Validating handleSteps Generator Behavior

The `handleSteps` function yields step instructions and tool calls in a specific order. Unit tests should drive the generator manually and assert on the yielded values:

```typescript
import { handleSteps } from '../thinker/thinker'

describe('thinker agent handleSteps', () => {
  test('yields STEP as first instruction', () => {
    const generator = handleSteps({ userMessage: 'Refactor this code' })
    const firstStep = generator.next()
    
    expect(firstStep.value).toBe('STEP')
    expect(firstStep.done).toBe(false)
  })
  
  test('yields set_output tool call in later steps', () => {
    const generator = handleSteps({ userMessage: 'Analyze this' })
    
    // Advance through initial steps
    generator.next()
    const secondStep = generator.next({ toolResult: 'intermediate data' })
    
    expect(secondStep.value).toHaveProperty('toolName', 'set_output')
  })
})

```

This pattern allows you to verify that the agent correctly sequences its internal logic, yields the expected step types, and formats tool calls appropriately before the code ever touches the runtime engine.

## Integration and End-to-End Testing

While unit tests verify internal logic, integration tests confirm that agents function correctly within the full Codebuff runtime environment, including tool execution and message handling.

### Running Full Agent Lifecycles with the SDK

Integration tests use the Codebuff SDK to spawn agents and execute complete workflows. These tests typically run against a test runtime or mocked LLM responses:

```typescript
import { run } from 'codebuff/sdk'

describe('editor agent integration', () => {
  test('completes full edit workflow', async () => {
    const result = await run({
      agentId: 'editor',
      userMessage: 'Add error handling to src/utils.ts',
      context: { files: ['src/utils.ts'] }
    })
    
    expect(result.output).toContain('try')
    expect(result.toolCalls).toContainEqual(
      expect.objectContaining({ toolName: 'write_file' })
    )
  })
})

```

These tests verify that the agent correctly interprets user messages, selects appropriate tools, and produces the expected final output.

### Testing Tool Interactions and Multi-Step Behavior

Complex agents like the *editor* and *thinker* require multiple steps and tool interactions. Integration tests should verify the complete sequence:

```typescript
describe('multi-step thinker agent', () => {
  test('reads files before generating output', async () => {
    const mockLLMResponses = [
      { toolCalls: [{ toolName: 'read_files', params: { paths: ['README.md'] } }] },
      { content: 'Analysis complete', toolCalls: [{ toolName: 'set_output', params: { value: 'result' } }] }
    ]
    
    const result = await runWithMockLLM('thinker', mockLLMResponses)
    
    expect(result.steps[0].toolCalls[0].toolName).toBe('read_files')
    expect(result.steps[1].toolCalls[0].toolName).toBe('set_output')
  })
})

```

These tests ensure that agents correctly chain tool calls, handle intermediate results, and reach completion states.

## Summary

- **Unit tests** validate the static `AgentDefinition` structure and `handleSteps` generator logic without launching the runtime, using patterns from [`agents/__tests__/editor.test.ts`](https://github.com/CodebuffAI/codebuff/blob/main/agents/__tests__/editor.test.ts).
- **Integration tests** execute full agent lifecycles via the SDK's `run` method, verifying tool interactions like `read_files`, `write_file`, and `set_output` in multi-step workflows.
- **Model mapping tests** ensure shorthand names like `'opus'` and `'gpt-5'` correctly resolve to full OpenRouter identifiers through the `createCodeEditor` factory.
- **Generator testing** allows direct invocation of `handleSteps` to assert on yield sequences, ensuring agents emit `'STEP'` instructions and tool calls in the correct order.

## Frequently Asked Questions

### What is the AgentDefinition interface in Codebuff?

The `AgentDefinition` interface, located in [`agents/types/agent-definition.ts`](https://github.com/CodebuffAI/codebuff/blob/main/agents/types/agent-definition.ts), is the typed contract that every Codebuff agent must implement. It defines required fields including `id`, `displayName`, `model`, `outputMode`, and the list of `toolNames`, ensuring that agents declare their capabilities and configuration explicitly before runtime execution.

### How do I test model shortcuts like 'opus' or 'gpt-5'?

Use the `createCodeEditor` factory function to instantiate agents with shorthand model names, then assert that the resulting `model` property maps to the correct full OpenRouter identifier. For example, verify that `createCodeEditor('opus').model` equals `'anthropic/claude-opus-4.6'` and that `'gpt-5'` resolves to `'openai/gpt-5'`.

### Should I mock the LLM when testing custom agents?

Yes, for **unit tests** you should mock or directly invoke the `handleSteps` generator to test logic without LLM dependencies. For **integration tests**, use mocked LLM responses that return predetermined tool calls and content, allowing you to verify that the agent correctly sequences operations like `read_files` followed by `set_output` without incurring API costs or latency.

### How do I test multi-step agent workflows?

Drive the `handleSteps` generator manually in unit tests to assert on each yielded value, verifying that the first yield is `'STEP'` and subsequent yields contain expected tool calls like `set_output`. For integration testing, use the SDK's `run` method with mocked LLM responses that simulate multiple turns, then assert that the agent correctly chains tool calls such as `read_files` before `write_file` and reaches the expected completion state.