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

> Learn to write effective unit and integration tests for your Custom Agent Service in SecureDesign. Discover mocking techniques and real environment testing strategies to ensure robust code.

- Repository: [Harold Martin/secure-design](https://github.com/hbmartin/secure-design)
- Tags: 
- Published: 2026-03-03

---

**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`](https://github.com/hbmartin/secure-design/blob/main/package.json) and [`tsconfig.test.json`](https://github.com/hbmartin/secure-design/blob/main/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`](https://github.com/hbmartin/secure-design/blob/main/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`](https://github.com/hbmartin/secure-design/blob/main/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`](https://github.com/hbmartin/secure-design/blob/main/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:

```typescript
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:

```typescript
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:

```typescript
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:

```typescript
// 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:

```typescript
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`](https://github.com/hbmartin/secure-design/blob/main/src/services/customAgentService.ts) | Core service implementation containing `setupWorkingDirectory`, tool creation, and streaming logic (lines 29‑119) |
| [`src/services/workspaceStateService.ts`](https://github.com/hbmartin/secure-design/blob/main/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`](https://github.com/hbmartin/secure-design/blob/main/package.json) | Defines test commands: `npm test`, `npm run test:agent`, `npm run test:core` |
| [`tsconfig.test.json`](https://github.com/hbmartin/secure-design/blob/main/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`](https://github.com/hbmartin/secure-design/blob/main/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`](https://github.com/hbmartin/secure-design/blob/main/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`](https://github.com/hbmartin/secure-design/blob/main/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.