ANTIGRAVITY LABJP
Articles/Agents & Manager
Agents & Manager/2026-04-13Intermediate

Unit Testing AgentKit 2.0 Agents with Vitest — A Practical Guide

Learn how to unit test and integration test MCP skill servers built with AgentKit 2.0 using Vitest. Covers mock strategies, InMemoryTransport integration tests, and CI/CD setup with real working code examples.

AgentKit17testing17VitestMCP17TypeScript11CI/CD16unit testing

Writing agent code is satisfying. Knowing it actually works under any condition? That's a different challenge. MCP skill servers exchange JSON over transports, call external APIs, and their behavior ultimately feeds into model decisions. Manually poking at them only takes you so far — every refactor becomes a gamble.

This guide uses the Weather skill server from the previous article as a concrete example, showing how to write unit tests and integration tests with Vitest. If you've been wondering "where do I even start testing agents?" — here's an answer in working code.

What Actually Makes Agent Testing Hard

Let's be honest about the challenges. MCP skill servers have three awkward properties:

1. Async + transport layer: The MCP SDK communicates via JSON-RPC over stdio or in-memory transports. There's no simple return value to assert against — process-level communication is involved.

2. External API dependencies: A weather skill hits OpenWeatherMap. A database skill hits a real DB. Running live API calls in every test is slow, expensive, and flaky on CI.

3. LLM non-determinism: The final agent output depends on model decisions. That part isn't testable — and it doesn't need to be.

This last point is the key insight. The right framing is:

┌─────────────────────────────────────┐
│  Test this (deterministic)           │
│  · Tool handler logic               │
│  · Input validation                 │
│  · Error handling & propagation     │
│  · External API calls via mocks     │
└─────────────────────────────────────┘
┌─────────────────────────────────────┐
│  Don't test this (non-deterministic) │
│  · Which tool the LLM selects       │
│  · Model-generated response text    │
└─────────────────────────────────────┘

With that separation clear, testing becomes tractable.

Project Structure and Vitest Setup

Adding a test environment to the Weather skill server from before:

weather-skill/
├── src/
   ├── index.ts          # MCP server entry point
   ├── tools/
   └── weather.ts    # Tool handler (test target)
   └── utils/
       └── openweather.ts # API client (mock target)
├── tests/
   ├── unit/
   └── weather.test.ts
   └── integration/
       └── server.test.ts
├── vitest.config.ts
└── package.json

Install Vitest:

npm install -D vitest @vitest/coverage-v8

Create vitest.config.ts:

import { defineConfig } from 'vitest/config'
 
export default defineConfig({
  test: {
    globals: true,
    environment: 'node',
    coverage: {
      provider: 'v8',
      reporter: ['text', 'lcov'],
      include: ['src/**/*.ts'],
      exclude: ['src/index.ts'],  // Skip entry point
    },
  },
})

Add scripts to package.json:

{
  "scripts": {
    "test": "vitest run",
    "test:watch": "vitest",
    "test:coverage": "vitest run --coverage"
  }
}

Unit Testing the Tool Handler

The Weather skill separates its API client into src/utils/openweather.ts and its handler logic into src/tools/weather.ts. That separation is what makes Vitest's vi.mock() practical.

The API client:

// src/utils/openweather.ts
export interface WeatherData {
  city: string
  temperature: number
  condition: string
  humidity: number
}
 
export async function fetchCurrentWeather(city: string): Promise<WeatherData> {
  const url = `https://api.openweathermap.org/data/2.5/weather?q=${city}&appid=${process.env.OPENWEATHER_API_KEY}&units=metric`
  const res = await fetch(url)
  if (\!res.ok) {
    throw new Error(`OpenWeatherMap API error: ${res.status}`)
  }
  const data = await res.json()
  return {
    city: data.name,
    temperature: Math.round(data.main.temp),
    condition: data.weather[0].description,
    humidity: data.main.humidity,
  }
}

The tool handler:

// src/tools/weather.ts
import { fetchCurrentWeather } from '../utils/openweather.js'
 
export async function handleGetWeather(args: { city: string }) {
  if (\!args.city || args.city.trim() === '') {
    throw new Error('city parameter is required')
  }
 
  const weather = await fetchCurrentWeather(args.city.trim())
  
  return {
    content: [
      {
        type: 'text' as const,
        text: `Current weather in ${weather.city}: ${weather.temperature}°C, ${weather.condition}, humidity ${weather.humidity}%`,
      },
    ],
  }
}

Now the unit tests:

// tests/unit/weather.test.ts
import { describe, it, expect, vi, beforeEach } from 'vitest'
import { handleGetWeather } from '../../src/tools/weather.js'
 
// Mock the entire openweather module
vi.mock('../../src/utils/openweather.js', () => ({
  fetchCurrentWeather: vi.fn(),
}))
 
import { fetchCurrentWeather } from '../../src/utils/openweather.js'
const mockFetchWeather = vi.mocked(fetchCurrentWeather)
 
describe('handleGetWeather', () => {
  beforeEach(() => {
    vi.clearAllMocks()
  })
 
  it('returns a text response for a valid city', async () => {
    mockFetchWeather.mockResolvedValue({
      city: 'Tokyo',
      temperature: 22,
      condition: 'clear sky',
      humidity: 60,
    })
 
    const result = await handleGetWeather({ city: 'Tokyo' })
 
    expect(result.content).toHaveLength(1)
    expect(result.content[0].type).toBe('text')
    expect(result.content[0].text).toContain('Tokyo')
    expect(result.content[0].text).toContain('22°C')
    expect(result.content[0].text).toContain('60%')
  })
 
  it('throws for an empty city string', async () => {
    await expect(handleGetWeather({ city: '' })).rejects.toThrow('city parameter is required')
    expect(mockFetchWeather).not.toHaveBeenCalled()
  })
 
  it('throws for a whitespace-only city string', async () => {
    await expect(handleGetWeather({ city: '   ' })).rejects.toThrow('city parameter is required')
  })
 
  it('propagates API errors', async () => {
    mockFetchWeather.mockRejectedValue(new Error('OpenWeatherMap API error: 404'))
    await expect(handleGetWeather({ city: 'InvalidCity' })).rejects.toThrow('OpenWeatherMap API error: 404')
  })
 
  it('trims whitespace before calling the API', async () => {
    mockFetchWeather.mockResolvedValue({
      city: 'Osaka',
      temperature: 25,
      condition: 'partly cloudy',
      humidity: 55,
    })
 
    await handleGetWeather({ city: '  Osaka  ' })
 
    expect(mockFetchWeather).toHaveBeenCalledWith('Osaka')
  })
})

Run npm test and all five tests should pass. This locks down the tool's behavior before you touch the MCP transport layer.

Testing the API Client — Mocking fetch

The fetchCurrentWeather function itself is worth testing. Here's how to mock the global fetch:

// tests/unit/openweather.test.ts
import { describe, it, expect, vi, beforeEach } from 'vitest'
import { fetchCurrentWeather } from '../../src/utils/openweather.js'
 
const mockFetch = vi.fn()
vi.stubGlobal('fetch', mockFetch)
 
describe('fetchCurrentWeather', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    process.env.OPENWEATHER_API_KEY = 'test-api-key'
  })
 
  it('parses and returns API response correctly', async () => {
    mockFetch.mockResolvedValue({
      ok: true,
      json: async () => ({
        name: 'Tokyo',
        main: { temp: 22.4, humidity: 60 },
        weather: [{ description: 'clear sky' }],
      }),
    })
 
    const result = await fetchCurrentWeather('Tokyo')
 
    expect(result).toEqual({
      city: 'Tokyo',
      temperature: 22,  // Math.round(22.4)
      condition: 'clear sky',
      humidity: 60,
    })
  })
 
  it('includes the API key and units in the request URL', async () => {
    mockFetch.mockResolvedValue({
      ok: true,
      json: async () => ({
        name: 'Tokyo',
        main: { temp: 20, humidity: 50 },
        weather: [{ description: 'clear' }],
      }),
    })
 
    await fetchCurrentWeather('Tokyo')
 
    const calledUrl = mockFetch.mock.calls[0][0] as string
    expect(calledUrl).toContain('appid=test-api-key')
    expect(calledUrl).toContain('units=metric')
    expect(calledUrl).toContain('q=Tokyo')
  })
 
  it('throws when the API responds with ok=false', async () => {
    mockFetch.mockResolvedValue({ ok: false, status: 404 })
    await expect(fetchCurrentWeather('NonExistentCity')).rejects.toThrow('OpenWeatherMap API error: 404')
  })
})

The Math.round test is easy to overlook but important — 22.4 becoming 22 is a specification, and without a test it's the kind of thing that silently breaks in a refactor.

Integration Testing the Full MCP Server

With unit tests covering the logic, integration tests verify the complete JSON-RPC flow. The MCP SDK provides InMemoryTransport for exactly this purpose — no need to spawn a real stdio process:

// tests/integration/server.test.ts
import { describe, it, expect, vi, beforeAll, afterAll } from 'vitest'
import { Server } from '@modelcontextprotocol/sdk/server/index.js'
import { Client } from '@modelcontextprotocol/sdk/client/index.js'
import { InMemoryTransport } from '@modelcontextprotocol/sdk/inMemory.js'
import { createWeatherServer } from '../../src/server.js'
 
// External API stays mocked in integration tests too
vi.mock('../../src/utils/openweather.js', () => ({
  fetchCurrentWeather: vi.fn().mockResolvedValue({
    city: 'Tokyo',
    temperature: 22,
    condition: 'clear sky',
    humidity: 60,
  }),
}))
 
describe('WeatherMCPServer integration tests', () => {
  let server: Server
  let client: Client
 
  beforeAll(async () => {
    server = createWeatherServer()
    client = new Client({ name: 'test-client', version: '1.0.0' }, { capabilities: {} })
 
    const [serverTransport, clientTransport] = InMemoryTransport.createLinkedPair()
    await Promise.all([
      server.connect(serverTransport),
      client.connect(clientTransport),
    ])
  })
 
  afterAll(async () => {
    await client.close()
    await server.close()
  })
 
  it('returns the tool list', async () => {
    const tools = await client.listTools()
    const toolNames = tools.tools.map(t => t.name)
    expect(toolNames).toContain('get_current_weather')
  })
 
  it('calls get_current_weather and returns a result', async () => {
    const result = await client.callTool({ name: 'get_current_weather', arguments: { city: 'Tokyo' } })
 
    expect(result.isError).toBeFalsy()
    expect(result.content).toHaveLength(1)
    expect(result.content[0].type).toBe('text')
    expect((result.content[0] as { type: string; text: string }).text).toContain('Tokyo')
  })
 
  it('throws for an unknown tool name', async () => {
    await expect(
      client.callTool({ name: 'nonexistent_tool', arguments: {} })
    ).rejects.toThrow()
  })
})

InMemoryTransport.createLinkedPair() creates two transport endpoints that are wired together in memory. The full JSON-RPC handshake — capability negotiation, listTools, callTool — runs exactly as it would in production, but without a network or process boundary.

Adding CI with GitHub Actions

Once tests are in place, automate them. Create .github/workflows/test.yml:

name: Test
 
on:
  push:
    branches: [main]
  pull_request:
    branches: [main]
 
jobs:
  test:
    runs-on: ubuntu-latest
 
    steps:
      - uses: actions/checkout@v4
 
      - uses: actions/setup-node@v4
        with:
          node-version: '20'
          cache: 'npm'
 
      - run: npm ci
 
      - name: Run tests with coverage
        run: npm run test:coverage
        env:
          OPENWEATHER_API_KEY: ${{ secrets.OPENWEATHER_API_KEY }}
          # Store a dummy value like "test-dummy-key" in secrets.
          # Tests mock the fetch call, so the real key isn't used.
 
      - name: Upload coverage
        uses: codecov/codecov-action@v4
        with:
          file: ./coverage/lcov.info
          fail_ci_if_error: false

One practical note: even though tests mock fetch, having process.env.OPENWEATHER_API_KEY referenced in source code means an undefined variable could cause issues in some environments. Store a dummy string like "test-dummy-key" in GitHub Secrets to keep things clean.

Three Things Worth Knowing Before You Start

After going through this for real, three pitfalls come up consistently.

vi.mock() is automatically hoisted to the top of the file. Vitest moves all vi.mock() calls before import statements at compile time. This means mocks are in place before any module runs. The practical implication: set dynamic mock values inside beforeEach using mockResolvedValue, not in the vi.mock() factory function, or per-test overrides won't work reliably.

InMemoryTransport API changed in SDK v1.9+. The createLinkedPair() static method was introduced then. If you're on an older version, you'll need to create two InMemoryTransport instances and manually link them. Check your package.json before copying the integration test above.

Always check result.isError before asserting on content. In MCP, tool errors don't throw — they return { isError: true, content: [...] }. If you skip the isError check and go straight to result.content[0].text, a failing tool call will produce a misleading test failure where the assertion passes on an error message rather than the expected output.

Where to Start

A three-stage approach works well in practice.

Start with just the tool handler unit tests. Mock external dependencies with vi.mock(), cover input validation and error propagation. Even this much makes refactoring significantly less stressful.

Add integration tests using InMemoryTransport. Verify the full JSON-RPC flow — listTools returns the right tools, callTool returns the right shape, unknown tools behave predictably.

Connect CI last. Once tests run on every PR, you can add new skills or refactor existing ones without worrying that something broke silently.

Build the skill in the previous article, add the tests from this one, and you have a production-grade MCP skill server worth shipping.

Share

Thank You for Reading

Antigravity Lab is ad-free, supported entirely by members like you. We publish practical guides daily with implementation code, benchmarks, and production-ready patterns. If you've found it useful, we'd love to have you on board.

  • Copy-paste ready implementation code
  • New advanced guides published daily
  • $5/mo or $10 for lifetime access
View Membership →

If you found this article helpful, a small tip ($1.50) would mean a lot to us. Your support helps keep this site ad-free and covers server and hosting costs.

Related Articles

Agents & Manager2026-07-05
Make the Self-Debugging Agent Walk the Logged-In and Post-Paywall Screens
By default, Antigravity 2.0's real-browser self-debug only sees the logged-out free view and reports success. To catch billing regressions, inject an authenticated session and paid state into the agent's browser and force coverage with assertions.
Agents & Manager2026-06-29
When Your Antigravity-Written Tests Only Look Green: Measuring Effectiveness with Mutation Testing
Tests written by an Antigravity agent can pass and still fail to catch the bug that matters. Here is how to measure their real effectiveness with mutation testing and only adopt tests after they kill the surviving mutants — with working code.
Agents & Manager2026-06-27
Pin Your Agent's Output With Golden Snapshots Before Switching Models
When Antigravity's engine moves to Gemini 3.5 Flash, an agent's output can drift silently. This walks through a golden-snapshot regression gate that catches the drift, with the actual test code and a migration-day checklist.
📚RECOMMENDED BOOKS
Build a Large Language Model (From Scratch)
Sebastian Raschka
LLM Dev
Prompt Engineering for LLMs
Berryman & Ziegler
Prompting
AI Engineering
Chip Huyen
AI Eng
* Contains affiliate links
See all →