How to Write Unit and Integration Tests for the Custom Agent Service in SecureDesign

Developers can test the Custom Agent Service by mocking VS Code APIs, the LLM provider, and file system calls in unit tests, or by using temporary workspaces and real file system operations in integration tests, leveraging the Mocha test harness defined in package.json and tsconfig.test.json.

The Custom Agent Service in the hbmartin/secure-design repository acts as the core bridge between the VS Code extension and the LLM SDK, handling tool execution, streaming responses, and workspace initialization. Writing unit and integration tests for custom agent service functionalities ensures that tool calls, error handling, and streaming logic remain stable as the codebase evolves. This guide demonstrates how to leverage the existing Mocha test harness to validate both isolated logic and end-to-end workflows.

Understanding the Custom Agent Service Architecture

Before writing tests, you must understand the components that require validation. The service implementation in src/services/customAgentService.ts contains several critical integration points:

• setupWorkingDirectory (lines 60‑118) initializes a .superdesign folder in the workspace or falls back to a temporary directory
• Tool creation (lines 29‑41) instantiates the tool set (read, write, edit, glob, grep, ls, bash, generateTheme) bound to an ExecutionContext
• Streaming logic (lines 63‑119) consumes ai.sdk.streamText and processes chunks via a switch statement handling text, tool-call, tool-result, and error states
• Error reporting via extractErrorMessage (lines 23‑44) and the centralized getLogger

Because the service relies on VS Code APIs (vscode.workspace.fs, ExtensionContext.secrets) and the LLM provider (@ai-sdk/provider), tests must isolate or simulate these dependencies.

Testing Strategy and Harness

The repository ships with a conventional VS Code extension test setup:

• Mocha test runner installed via @vscode/test-electron and @vscode/test-cli
• Test scripts in package.json: "test": "vscode-test" for standard suites and "test:agent": "tsc --project tsconfig.test.json && node dist-test/test/llm-service.test.js" for agent-specific validation
• Test directory configured in tsconfig.test.json to include src/test/**/*

You can execute tests via npm test for general validation or npm run test:agent for agent-specific flows that require the specialized TypeScript configuration.

Unit Tests vs. Integration Tests

Unit tests stub all external collaborators—file system, secret storage, and LLM streams—to validate the internal state machine of the switch handling logic. These run fast and suit the npm run test:core suite.

Integration tests retain real file-system access using temporary workspace folders and may use the actual LLM provider (if a test-only API key is available) or a local mock server. These validate end-to-end behavior and run via npm run test:agent inside a headless VS Code extension host.

Writing Unit Tests for Custom Agent Service

Unit testing requires mocking three layers: VS Code APIs, the workspace state service, and the LLM stream.

Mocking VS Code APIs and File System

Use sinon to stub vscode.workspace.fs methods and prevent disk side effects:

import * as sinon from 'sinon';
import * as vscode from 'vscode';

const fakeFs = {
  stat: sinon.stub().resolves({} as any),
  createDirectory: sinon.stub().resolves(),
  readFile: sinon.stub().resolves(new Uint8Array()),
  writeFile: sinon.stub().resolves(),
};

sinon.replace(vscode.workspace, 'fs', fakeFs as any);
sinon.stub(vscode.window, 'showWarningMessage').resolves();

Stub vscode.SecretStorage to isolate credential operations:

const fakeSecret = {
  get: sinon.stub().resolves(undefined),
  store: sinon.stub().resolves(),
  delete: sinon.stub().resolves(),
};

sinon.replace(vscode, 'SecretStorage', fakeSecret as any);

Simulating LLM Stream Responses

Create a fake LanguageModelV2 that yields deterministic chunks matching the protocol expected by the service's streaming logic:

import type { LanguageModelV2 } from '@ai-sdk/provider';

const fakeModel: LanguageModelV2 = {
  async *streamText(): AsyncIterable<any> {
    // Simulate a tool call to the write tool
    yield { type: 'tool-call', toolCallId: 'c1', toolName: 'write', input: { file_path: 'design.html', content: '<h1>Hello</h1>' } };
    // Return the tool result
    yield { type: 'tool-result', toolCallId: 'c1', toolName: 'write', output: 'ok' };
    // Signal completion
    yield { type: 'finish', finishReason: 'stop' };
  },
} as any;

Complete Unit Test Example

The following test validates that a write tool call results in a tool-result message and invokes the VS Code file system API:

// src/test/agent/customAgentService.test.ts
import * as assert from 'assert';
import * as sinon from 'sinon';
import * as vscode from 'vscode';
import { CustomAgentService } from '../../src/services/customAgentService';
import { WorkspaceStateService } from '../../src/services/workspaceStateService';

// Setup mocks as described above
const fakeFs = {
  stat: sinon.stub().resolves({} as any),
  createDirectory: sinon.stub().resolves(),
  readFile: sinon.stub().resolves(new Uint8Array()),
  writeFile: sinon.stub().resolves(),
};
sinon.replace(vscode.workspace, 'fs', fakeFs as any);

const fakeSecret = {
  get: sinon.stub().resolves(undefined),
  store: sinon.stub().resolves(),
  delete: sinon.stub().resolves(),
};
sinon.replace(vscode, 'SecretStorage', fakeSecret as any);

// Initialize WorkspaceStateService with minimal context
const wsService = WorkspaceStateService.getInstance();
wsService.initialize({} as vscode.ExtensionContext);

// Instantiate service
const agent = new CustomAgentService(wsService);

test('CustomAgentService processes a write‑tool call', async () => {
  const history = [{ role: 'user' as const, content: 'Create a file' }];
  const abort = new AbortController();
  
  // Execute query (callback can be no-op for unit tests)
  const result = await agent.query(history, abort, () => {});
  
  // Assert final message structure
  const last = result[result.length - 1];
  assert.strictEqual(last.role, 'tool');
  const toolPart = (last.content as any)[0];
  assert.strictEqual(toolPart.type, 'tool-result');
  assert.strictEqual(toolPart.toolName, 'write');
  assert.strictEqual(toolPart.output, 'ok');
  
  // Verify file system interaction
  sinon.assert.calledOnce(fakeFs.writeFile);
});

Writing Integration Tests for End-to-End Validation

Integration tests verify the service against real file system operations. Use a temporary directory and point the workspace to it:

import * as tmp from 'tmp-promise';
import * as vscode from 'vscode';
import * as sinon from 'sinon';

// Inside your test case:
const { path: tmpDir } = await tmp.dir({ unsafeCleanup: true });
const workspaceUri = vscode.Uri.file(tmpDir);

// Redirect workspace folders to temp location
sinon.stub(vscode.workspace, 'workspaceFolders').value([{ uri: workspaceUri }]);

// Now instantiate CustomAgentService and run agent.query
// The service will create a real `.superdesign` folder inside tmpDir

Run these tests with npm run test:agent, which compiles the test suite via tsc --project tsconfig.test.json and executes it inside the headless VS Code extension host exactly as CI does.

Key Files for Test Implementation

File	Purpose
src/services/customAgentService.ts	Core service implementation containing setupWorkingDirectory, tool creation, and streaming logic (lines 29‑119)
src/services/workspaceStateService.ts	Provides secret storage and workspace state required by the agent
src/tools/*.ts	Individual tool implementations (read, write, edit, etc.) exercised during testing
package.json	Defines test commands: npm test, npm run test:agent, npm run test:core
tsconfig.test.json	Configures TypeScript compilation for the src/test/**/* directory

Summary

• Mock external dependencies—VS Code APIs (vscode.workspace.fs, SecretStorage) and the LLM provider—to enable fast, deterministic unit tests
• Use temporary directories and real file system operations for integration tests that validate setupWorkingDirectory behavior
• Simulate LLM streams by implementing fake streamText iterators that yield tool-call, tool-result, and finish chunks
• Leverage the existing harness via npm run test:agent to run compiled tests in a headless VS Code extension host
• Target specific logic such as the chunk-handling switch statement (lines 63‑119) and error extraction (lines 23‑44) for comprehensive coverage

Frequently Asked Questions

How do I mock the LLM stream for deterministic testing?

Create a fake object implementing LanguageModelV2 with an async *streamText() generator method. Yield objects with type: 'tool-call', type: 'tool-result', and type: 'finish' properties to simulate the exact sequence the service expects. This approach tests the switch logic in src/services/customAgentService.ts without network calls.

Can I run integration tests without a real LLM API key?

Yes. Integration tests can focus on the file-system and workspace initialization layers while stubbing the LLM provider, or you can point the service to a local mock server that implements the streamText protocol. The test harness in package.json supports both configurations.

What is the difference between npm test and npm run test:agent?

npm test executes the standard VS Code extension test suite via vscode-test, typically running unit tests with extensive mocking. npm run test:agent explicitly compiles TypeScript from tsconfig.test.json and executes the compiled JavaScript in Node, often used for agent-specific integration tests that require the full service initialization sequence.

How do I verify that a tool actually modified the file system?

In unit tests, assert that sinon stubs like fakeFs.writeFile were called with expected arguments. In integration tests, use Node.js fs promises to read the temporary workspace directory after agent.query resolves, verifying that files exist at the paths specified in the tool calls.

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