# How to Write Unit and Integration Tests for Ax Code: Vitest Patterns and Examples

> Learn to write unit and integration tests for Ax code using Vitest patterns. Discover deterministic LLM mocks with AxMockAIService for efficient testing, reducing API dependencies.

- Repository: [Ax/ax](https://github.com/ax-llm/ax)
- Tags: how-to-guide
- Published: 2026-02-25

---

**Ax uses Vitest for testing, co-locating `*.test.ts` unit files and `*.integration.test.ts` files alongside source code, with `AxMockAIService` providing deterministic LLM mocks that eliminate external API dependencies.**

The ax-llm/ax repository is a TypeScript framework for building LLM-powered agents. Writing unit and integration tests for Ax code ensures your agent logic, signature derivations, and runtime integrations behave deterministically without consuming costly API tokens. The test suite relies on Vitest and follows a factory-function API style (`agent()`, `ax()`, `ai()`) rather than deprecated constructors.

## Test Setup Conventions

Ax distinguishes between unit and integration tests through file naming and mock boundaries:

| Test type | File pattern | Purpose |
|-----------|--------------|---------|
| **Unit** | `*.test.ts` | Verify a single module’s behavior in isolation (e.g., signature derivation). |
| **Integration** | `*.integration.test.ts` or standard tests with mocks | Exercise multiple modules together (e.g., the full actor/responder loop). |

All tests reside in `src/` alongside source files rather than a separate `tests/` directory.

### Core Testing Primitives

- **`AxMockAIService`** – Simulates LLM responses via `new AxMockAIService({ features, chatResponse })`. Supports static responses, functions, or streaming `ReadableStream` objects.
- **Runtime stubs** – Minimal `AxCodeRuntime` implementations that only implement `execute`, returning `'ok'` or evaluating simple JS.
- **Fluent signature API** – Use `s('input:string -> output:string')` for type-safe signatures that tests can introspect via `sig.getInputFields()`.
- **Agent factory** – Construct agents with `agent('input:string -> output:string', { runtime, contextFields, ai })` to centralize configuration.

## Unit Test Patterns

Unit tests target discrete modules like `AxAgent` or `AxGen` without invoking external services.

### Validating Constructor Logic

Verify that `AxAgent` throws appropriate errors when misconfigured:

```typescript
import { AxAgent } from './agent.js';
import { defaultRlmFields } from './agent.test.js';

describe('AxAgent', () => {
  it('throws when getFunction() is called without an identity', () => {
    const a = new AxAgent(
      { signature: 'userQuery:string -> agentResponse:string' },
      { ...defaultRlmFields }
    );
    expect(() => a.getFunction()).toThrow(/agentIdentity/);
  });
});

```

*Reference*: [src/ax/prompts/agent.test.ts → line 51‑61](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/agent.test.ts#L51-L61)

### Testing Signature Derivation

Assert that the agent correctly manipulates input/output fields when constructing sub-agents:

```typescript
const testAgent = agent('context:string, query:string -> answer:string', {
  contextFields: ['context'],
  runtime,
});
const actorSig = (testAgent as any).actorProgram.getSignature();
expect(actorSig.getInputFields().some(f => f.name === 'actionLog')).toBeTruthy();
expect(actorSig.getOutputFields()[0].name).toBe('javascriptCode');

```

*Reference*: [src/ax/prompts/agent.test.ts → line 59‑81](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/agent.test.ts#L59-L81)

### Streaming vs Non-Streaming Forward

Test both modalities using `AxMockAIService` with feature flags:

```typescript
const ai = new AxMockAIService({ 
  features: { streaming: true }, 
  chatResponse: streamingResponse 
});
const gen = new AxGen<{ q: string }, { a: string }>('q:string -> a:string');
const result = await gen.forward(ai, { q: 'test' }, { stream: true });
expect(result.a).toContain('chunk 1');

```

*Reference*: [src/ax/dsp/generate.test.ts → line 71‑40](https://github.com/ax-llm/ax/blob/main/src/ax/dsp/generate.test.ts#L71-L40)

## Integration Test Patterns

Integration tests verify multi-module workflows such as the full agent loop or runtime execution.

### Full Actor/Responder Loop

Simulate a complete agent execution where the mock AI returns intermediate code, then a final answer:

```typescript
const testMockAI = new AxMockAIService({
  features: { functions: false, streaming: false },
  chatResponse: async req => {
    const prompt = String(req.chatPrompt[0]?.content ?? '');
    if (prompt.includes('Code Generation Agent')) {
      // Return code on first call, final() on second
      return { results: [{ content: 'const x = 42;' }], modelUsage };
    }
    if (prompt.includes('Answer Synthesis Agent')) {
      return { results: [{ content: 'The answer is 42' }], modelUsage };
    }
    return { results: [{ content: 'fallback' }], modelUsage };
  },
});

const runtime: AxCodeRuntime = { 
  execute: async (code) => eval(code) 
};

const testAgent = agent('context:string, query:string -> answer:string', {
  ai: testMockAI,
  contextFields: ['context'],
  runtime,
});

const result = await testAgent.forward(testMockAI, { context: 'x', query: 'y' });
expect(result.answer).toBe('The answer is 42');

```

*Reference*: [src/ax/prompts/agent.test.ts → lines 44‑58](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/agent.test.ts#L44-L58)

### JavaScript Runtime End-to-End

Test the sandbox execution layer directly:

```typescript
describe('AxJSRuntime integration', () => {
  it('executes a simple script and returns the result', async () => {
    const runtime = await import('../funcs/jsRuntime.js');
    const result = await runtime.evaluate('1 + 2');
    expect(result).toBe(3);
  });
});

```

*Reference*: [src/ax/funcs/jsRuntime.integration.test.ts](https://github.com/ax-llm/ax/blob/main/src/ax/funcs/jsRuntime.integration.test.ts)

## Common Helper Utilities

Reusing these exported helpers from existing test files keeps new tests concise and behaviorally consistent:

| File | Export | Purpose |
|------|--------|---------|
| [`src/ax/prompts/agent.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/agent.test.ts) | `defaultRuntime`, `defaultRlmFields` | Minimal runtime for tests that don’t need code execution. |
| [`src/ax/dsp/generate.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/dsp/generate.test.ts) | `createStreamingResponse` | Generates a mock `ReadableStream` of LLM chunks for streaming tests. |
| [`src/ax/util/apicall.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/util/apicall.test.ts) | `makeModelUsage` | Generates a stubbed `modelUsage` object for AI-service mocks. |

## Running the Tests

Execute the suite using the npm scripts defined in the workspace [`package.json`](https://github.com/ax-llm/ax/blob/main/package.json):

```bash

# Unit tests only (fast, no external calls)

npm run test:unit

# All workspace tests (unit + integration)

npm run test

```

These scripts invoke **Vitest** with the `run-s test:*` pattern, ensuring consistent CI and local behavior.

## Key Files to Study

Study these reference implementations to understand advanced patterns:

- **[`src/ax/prompts/agent.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/prompts/agent.test.ts)** – Comprehensive tests for `AxAgent` covering signature derivation, description injection, and the full actor/responder loop.
- **[`src/ax/dsp/generate.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/dsp/generate.test.ts)** – Tests the `AxGen` generator’s `forward` method, covering non-streaming, streaming aggregation, and tool-function stop handling.
- **[`src/ax/funcs/jsRuntime.integration.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/funcs/jsRuntime.integration.test.ts)** – End-to-end integration of the JavaScript runtime sandbox.
- **[`src/ax/flow/flow.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/flow/flow.test.ts)** – Demonstrates integration with the broader ecosystem, parallel execution, and caching for larger-scale flow tests.

## Summary

- **Co-locate tests** with source files using `*.test.ts` for unit tests and `*.integration.test.ts` for multi-module integration tests.
- **Use `AxMockAIService`** to simulate LLM responses deterministically, toggling `features.streaming` and `features.functions` as needed.
- **Prefer factory functions** (`agent()`, `ax()`) over constructors to align with the modern API surface.
- **Import shared helpers** like `defaultRuntime` and `makeModelUsage` from existing test files to avoid duplication.
- **Run `npm run test:unit`** for rapid feedback during development, reserving `npm run test` for full CI validation.

## Frequently Asked Questions

### What testing framework does Ax use?

Ax uses **Vitest** for all unit and integration tests. The framework runs alongside TypeScript source files in `src/` and supports both standard [`.test.ts`](https://github.com/ax-llm/ax/blob/main/.test.ts) files and [`.integration.test.ts`](https://github.com/ax-llm/ax/blob/main/.integration.test.ts) variants for different test scopes.

### How do I mock LLM responses in Ax tests?

Instantiate **`AxMockAIService`** with a `chatResponse` function or static object. Configure `features` to declare capabilities (streaming, functions) so the Ax internals route correctly. Return `ReadableStream` instances for streaming tests or plain objects for standard generation.

### Where should I place integration tests in the Ax repository?

Place integration tests either in `src/<module>/**/*.integration.test.ts` files or within standard test files that spin up multiple real modules (e.g., an agent plus a mock AI service). The [`src/ax/funcs/jsRuntime.integration.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/funcs/jsRuntime.integration.test.ts) file demonstrates the dedicated integration file pattern.

### How do I test streaming responses in Ax?

Pass `features: { streaming: true }` to `AxMockAIService`, then invoke `forward` with `{ stream: true }`. Use the `createStreamingResponse` helper from [`src/ax/dsp/generate.test.ts`](https://github.com/ax-llm/ax/blob/main/src/ax/dsp/generate.test.ts) to generate realistic chunk streams, and assert that the final aggregated result contains expected partial content.