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

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:

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

Testing Signature Derivation

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

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

Streaming vs Non-Streaming Forward

Test both modalities using AxMockAIService with feature flags:

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

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:

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

JavaScript Runtime End-to-End

Test the sandbox execution layer directly:

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

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	defaultRuntime, defaultRlmFields	Minimal runtime for tests that don’t need code execution.
src/ax/dsp/generate.test.ts	createStreamingResponse	Generates a mock ReadableStream of LLM chunks for streaming tests.
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:

# 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 – Comprehensive tests for AxAgent covering signature derivation, description injection, and the full actor/responder loop.
• 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 – End-to-end integration of the JavaScript runtime sandbox.
• 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 files and .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 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 to generate realistic chunk streams, and assert that the final aggregated result contains expected partial content.

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