5 files changed, 609 insertions, 179 deletions
diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
new file mode 100644
index 00000000..80949203
--- /dev/null
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -0,0 +1,92 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi } from 'vitest';
+import { webLoginClient } from './oauth2.js';
+import { OAuth2Client } from 'google-auth-library';
+import http from 'http';
+import open from 'open';
+import crypto from 'crypto';
+
+vi.mock('google-auth-library');
+vi.mock('http');
+vi.mock('open');
+vi.mock('crypto');
+
+describe('oauth2', () => {
+  it('should perform a web login', async () => {
+    const mockAuthUrl = 'https://example.com/auth';
+    const mockCode = 'test-code';
+    const mockState = 'test-state';
+    const mockTokens = {
+      access_token: 'test-access-token',
+      refresh_token: 'test-refresh-token',
+    };
+
+    const mockGenerateAuthUrl = vi.fn().mockReturnValue(mockAuthUrl);
+    const mockGetToken = vi.fn().mockResolvedValue({ tokens: mockTokens });
+    const mockSetCredentials = vi.fn();
+    const mockOAuth2Client = {
+      generateAuthUrl: mockGenerateAuthUrl,
+      getToken: mockGetToken,
+      setCredentials: mockSetCredentials,
+    } as unknown as OAuth2Client;
+    vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client);
+
+    vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never);
+    vi.mocked(open).mockImplementation(async () => ({}) as never);
+
+    let requestCallback!: (
+      req: http.IncomingMessage,
+      res: http.ServerResponse,
+    ) => void;
+    const mockHttpServer = {
+      listen: vi.fn((port: number, callback?: () => void) => {
+        if (callback) {
+          callback();
+        }
+      }),
+      close: vi.fn((callback?: () => void) => {
+        if (callback) {
+          callback();
+        }
+      }),
+      on: vi.fn(),
+      address: () => ({ port: 1234 }),
+    };
+    vi.mocked(http.createServer).mockImplementation((cb) => {
+      requestCallback = cb as (
+        req: http.IncomingMessage,
+        res: http.ServerResponse,
+      ) => void;
+      return mockHttpServer as unknown as http.Server;
+    });
+
+    const clientPromise = webLoginClient();
+
+    // Wait for the server to be created
+    await new Promise((resolve) => setTimeout(resolve, 0));
+
+    const mockReq = {
+      url: `/oauth2callback?code=${mockCode}&state=${mockState}`,
+    } as http.IncomingMessage;
+    const mockRes = {
+      writeHead: vi.fn(),
+      end: vi.fn(),
+    } as unknown as http.ServerResponse;
+
+    if (requestCallback) {
+      await requestCallback(mockReq, mockRes);
+    }
+
+    const client = await clientPromise;
+
+    expect(open).toHaveBeenCalledWith(mockAuthUrl);
+    expect(mockGetToken).toHaveBeenCalledWith(mockCode);
+    expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
+    expect(client).toBe(mockOAuth2Client);
+  });
+});
diff --git a/packages/core/src/code_assist/server.test.ts b/packages/core/src/code_assist/server.test.ts
new file mode 100644
index 00000000..bb033163
--- /dev/null
+++ b/packages/core/src/code_assist/server.test.ts
@@ -0,0 +1,151 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi } from 'vitest';
+import { CodeAssistServer } from './server.js';
+import { OAuth2Client } from 'google-auth-library';
+
+vi.mock('google-auth-library');
+
+describe('CodeAssistServer', () => {
+  it('should be able to be constructed', () => {
+    const auth = new OAuth2Client();
+    const server = new CodeAssistServer(auth, 'test-project');
+    expect(server).toBeInstanceOf(CodeAssistServer);
+  });
+
+  it('should call the generateContent endpoint', async () => {
+    const auth = new OAuth2Client();
+    const server = new CodeAssistServer(auth, 'test-project');
+    const mockResponse = {
+      response: {
+        candidates: [
+          {
+            index: 0,
+            content: {
+              role: 'model',
+              parts: [{ text: 'response' }],
+            },
+            finishReason: 'STOP',
+            safetyRatings: [],
+          },
+        ],
+      },
+    };
+    vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse);
+
+    const response = await server.generateContent({
+      model: 'test-model',
+      contents: [{ role: 'user', parts: [{ text: 'request' }] }],
+    });
+
+    expect(server.callEndpoint).toHaveBeenCalledWith(
+      'generateContent',
+      expect.any(Object),
+    );
+    expect(response.candidates?.[0]?.content?.parts?.[0]?.text).toBe(
+      'response',
+    );
+  });
+
+  it('should call the generateContentStream endpoint', async () => {
+    const auth = new OAuth2Client();
+    const server = new CodeAssistServer(auth, 'test-project');
+    const mockResponse = (async function* () {
+      yield {
+        response: {
+          candidates: [
+            {
+              index: 0,
+              content: {
+                role: 'model',
+                parts: [{ text: 'response' }],
+              },
+              finishReason: 'STOP',
+              safetyRatings: [],
+            },
+          ],
+        },
+      };
+    })();
+    vi.spyOn(server, 'streamEndpoint').mockResolvedValue(mockResponse);
+
+    const stream = await server.generateContentStream({
+      model: 'test-model',
+      contents: [{ role: 'user', parts: [{ text: 'request' }] }],
+    });
+
+    for await (const res of stream) {
+      expect(server.streamEndpoint).toHaveBeenCalledWith(
+        'streamGenerateContent',
+        expect.any(Object),
+      );
+      expect(res.candidates?.[0]?.content?.parts?.[0]?.text).toBe('response');
+    }
+  });
+
+  it('should call the onboardUser endpoint', async () => {
+    const auth = new OAuth2Client();
+    const server = new CodeAssistServer(auth, 'test-project');
+    const mockResponse = {
+      name: 'operations/123',
+      done: true,
+    };
+    vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse);
+
+    const response = await server.onboardUser({
+      tierId: 'test-tier',
+      cloudaicompanionProject: 'test-project',
+      metadata: {},
+    });
+
+    expect(server.callEndpoint).toHaveBeenCalledWith(
+      'onboardUser',
+      expect.any(Object),
+    );
+    expect(response.name).toBe('operations/123');
+  });
+
+  it('should call the loadCodeAssist endpoint', async () => {
+    const auth = new OAuth2Client();
+    const server = new CodeAssistServer(auth, 'test-project');
+    const mockResponse = {
+      // TODO: Add mock response
+    };
+    vi.spyOn(server, 'callEndpoint').mockResolvedValue(mockResponse);
+
+    const response = await server.loadCodeAssist({
+      metadata: {},
+    });
+
+    expect(server.callEndpoint).toHaveBeenCalledWith(
+      'loadCodeAssist',
+      expect.any(Object),
+    );
+    expect(response).toBe(mockResponse);
+  });
+
+  it('should return 0 for countTokens', async () => {
+    const auth = new OAuth2Client();
+    const server = new CodeAssistServer(auth, 'test-project');
+    const response = await server.countTokens({
+      model: 'test-model',
+      contents: [{ role: 'user', parts: [{ text: 'request' }] }],
+    });
+    expect(response.totalTokens).toBe(0);
+  });
+
+  it('should throw an error for embedContent', async () => {
+    const auth = new OAuth2Client();
+    const server = new CodeAssistServer(auth, 'test-project');
+    await expect(
+      server.embedContent({
+        model: 'test-model',
+        contents: [{ role: 'user', parts: [{ text: 'request' }] }],
+      }),
+    ).rejects.toThrow();
+  });
+});
diff --git a/packages/core/src/core/contentGenerator.test.ts b/packages/core/src/core/contentGenerator.test.ts
new file mode 100644
index 00000000..8716a441
--- /dev/null
+++ b/packages/core/src/core/contentGenerator.test.ts
@@ -0,0 +1,49 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi } from 'vitest';
+import { createContentGenerator } from './contentGenerator.js';
+import { createCodeAssistContentGenerator } from '../code_assist/codeAssist.js';
+import { GoogleGenAI } from '@google/genai';
+
+vi.mock('../code_assist/codeAssist.js');
+vi.mock('@google/genai');
+
+describe('contentGenerator', () => {
+  it('should create a CodeAssistContentGenerator', async () => {
+    const mockGenerator = {} as unknown;
+    vi.mocked(createCodeAssistContentGenerator).mockResolvedValue(
+      mockGenerator as never,
+    );
+    const generator = await createContentGenerator({
+      model: 'test-model',
+      codeAssist: true,
+    });
+    expect(createCodeAssistContentGenerator).toHaveBeenCalled();
+    expect(generator).toBe(mockGenerator);
+  });
+
+  it('should create a GoogleGenAI content generator', async () => {
+    const mockGenerator = {
+      models: {},
+    } as unknown;
+    vi.mocked(GoogleGenAI).mockImplementation(() => mockGenerator as never);
+    const generator = await createContentGenerator({
+      model: 'test-model',
+      apiKey: 'test-api-key',
+    });
+    expect(GoogleGenAI).toHaveBeenCalledWith({
+      apiKey: 'test-api-key',
+      vertexai: undefined,
+      httpOptions: {
+        headers: {
+          'User-Agent': expect.any(String),
+        },
+      },
+    });
+    expect(generator).toBe((mockGenerator as GoogleGenAI).models);
+  });
+});
diff --git a/packages/core/src/telemetry/telemetry.test.ts b/packages/core/src/telemetry/telemetry.test.ts
new file mode 100644
index 00000000..7cd6f95f
--- /dev/null
+++ b/packages/core/src/telemetry/telemetry.test.ts
@@ -0,0 +1,72 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import {
+  initializeTelemetry,
+  shutdownTelemetry,
+  isTelemetrySdkInitialized,
+} from './sdk.js';
+import { Config } from '../config/config.js';
+import { NodeSDK } from '@opentelemetry/sdk-node';
+import * as loggers from './loggers.js';
+
+vi.mock('@opentelemetry/sdk-node');
+vi.mock('../config/config.js');
+vi.mock('./loggers.js');
+
+describe('telemetry', () => {
+  let mockConfig: Config;
+  let mockNodeSdk: NodeSDK;
+
+  beforeEach(() => {
+    vi.resetAllMocks();
+
+    mockConfig = new Config({
+      sessionId: 'test-session-id',
+      contentGeneratorConfig: {
+        model: 'test-model',
+      },
+      targetDir: '/test/dir',
+      debugMode: false,
+      cwd: '/test/dir',
+    });
+    vi.spyOn(mockConfig, 'getTelemetryEnabled').mockReturnValue(true);
+    vi.spyOn(mockConfig, 'getTelemetryOtlpEndpoint').mockReturnValue(
+      'http://localhost:4317',
+    );
+    vi.spyOn(mockConfig, 'getSessionId').mockReturnValue('test-session-id');
+    vi.spyOn(loggers, 'logCliConfiguration').mockImplementation(() => {});
+
+    mockNodeSdk = {
+      start: vi.fn(),
+      shutdown: vi.fn().mockResolvedValue(undefined),
+    } as unknown as NodeSDK;
+    vi.mocked(NodeSDK).mockImplementation(() => mockNodeSdk);
+  });
+
+  afterEach(async () => {
+    // Ensure we shut down telemetry even if a test fails.
+    if (isTelemetrySdkInitialized()) {
+      await shutdownTelemetry();
+    }
+  });
+
+  it('should initialize the telemetry service', () => {
+    initializeTelemetry(mockConfig);
+
+    expect(NodeSDK).toHaveBeenCalled();
+    expect(mockNodeSdk.start).toHaveBeenCalled();
+    expect(loggers.logCliConfiguration).toHaveBeenCalledWith(mockConfig);
+  });
+
+  it('should shutdown the telemetry service', async () => {
+    initializeTelemetry(mockConfig);
+    await shutdownTelemetry();
+
+    expect(mockNodeSdk.shutdown).toHaveBeenCalled();
+  });
+});
diff --git a/packages/core/src/utils/editor.test.ts b/packages/core/src/utils/editor.test.ts
index 8bc7a49c..84cca8b8 100644
--- a/packages/core/src/utils/editor.test.ts
+++ b/packages/core/src/utils/editor.test.ts
@@ -19,6 +19,7 @@ import {
   openDiff,
   allowEditorTypeInSandbox,
   isEditorAvailable,
+  type EditorType,
 } from './editor.js';
 import { execSync, spawn } from 'child_process';
 
@@ -27,225 +28,290 @@ vi.mock('child_process', () => ({
   spawn: vi.fn(),
 }));
 
-describe('checkHasEditorType', () => {
+const originalPlatform = process.platform;
+
+describe('editor utils', () => {
   beforeEach(() => {
     vi.clearAllMocks();
-  });
-
-  it('should return true for vscode if "code" command exists', () => {
-    (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/code'));
-    expect(checkHasEditorType('vscode')).toBe(true);
-    const expectedCommand =
-      process.platform === 'win32' ? 'where.exe code.cmd' : 'command -v code';
-    expect(execSync).toHaveBeenCalledWith(expectedCommand, {
-      stdio: 'ignore',
+    delete process.env.SANDBOX;
+    Object.defineProperty(process, 'platform', {
+      value: originalPlatform,
+      writable: true,
     });
   });
 
-  it('should return false for vscode if "code" command does not exist', () => {
-    (execSync as Mock).mockImplementation(() => {
-      throw new Error();
+  afterEach(() => {
+    vi.restoreAllMocks();
+    delete process.env.SANDBOX;
+    Object.defineProperty(process, 'platform', {
+      value: originalPlatform,
+      writable: true,
     });
-    expect(checkHasEditorType('vscode')).toBe(false);
   });
 
-  it('should return true for windsurf if "windsurf" command exists', () => {
-    (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/windsurf'));
-    expect(checkHasEditorType('windsurf')).toBe(true);
-    expect(execSync).toHaveBeenCalledWith('command -v windsurf', {
-      stdio: 'ignore',
-    });
-  });
+  describe('checkHasEditorType', () => {
+    const testCases: Array<{
+      editor: EditorType;
+      command: string;
+      win32Command: string;
+    }> = [
+      { editor: 'vscode', command: 'code', win32Command: 'code.cmd' },
+      { editor: 'windsurf', command: 'windsurf', win32Command: 'windsurf' },
+      { editor: 'cursor', command: 'cursor', win32Command: 'cursor' },
+      { editor: 'vim', command: 'vim', win32Command: 'vim' },
+    ];
 
-  it('should return false for windsurf if "windsurf" command does not exist', () => {
-    (execSync as Mock).mockImplementation(() => {
-      throw new Error();
-    });
-    expect(checkHasEditorType('windsurf')).toBe(false);
-  });
+    for (const { editor, command, win32Command } of testCases) {
+      describe(`${editor}`, () => {
+        it(`should return true if "${command}" command exists on non-windows`, () => {
+          Object.defineProperty(process, 'platform', { value: 'linux' });
+          (execSync as Mock).mockReturnValue(
+            Buffer.from(`/usr/bin/${command}`),
+          );
+          expect(checkHasEditorType(editor)).toBe(true);
+          expect(execSync).toHaveBeenCalledWith(`command -v ${command}`, {
+            stdio: 'ignore',
+          });
+        });
 
-  it('should return true for cursor if "cursor" command exists', () => {
-    (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/cursor'));
-    expect(checkHasEditorType('cursor')).toBe(true);
-    expect(execSync).toHaveBeenCalledWith('command -v cursor', {
-      stdio: 'ignore',
-    });
-  });
+        it(`should return false if "${command}" command does not exist on non-windows`, () => {
+          Object.defineProperty(process, 'platform', { value: 'linux' });
+          (execSync as Mock).mockImplementation(() => {
+            throw new Error();
+          });
+          expect(checkHasEditorType(editor)).toBe(false);
+        });
 
-  it('should return false for cursor if "cursor" command does not exist', () => {
-    (execSync as Mock).mockImplementation(() => {
-      throw new Error();
-    });
-    expect(checkHasEditorType('cursor')).toBe(false);
-  });
+        it(`should return true if "${win32Command}" command exists on windows`, () => {
+          Object.defineProperty(process, 'platform', { value: 'win32' });
+          (execSync as Mock).mockReturnValue(
+            Buffer.from(`C:\\Program Files\\...\\${win32Command}`),
+          );
+          expect(checkHasEditorType(editor)).toBe(true);
+          expect(execSync).toHaveBeenCalledWith(`where.exe ${win32Command}`, {
+            stdio: 'ignore',
+          });
+        });
 
-  it('should return true for vim if "vim" command exists', () => {
-    (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/vim'));
-    expect(checkHasEditorType('vim')).toBe(true);
-    const expectedCommand =
-      process.platform === 'win32' ? 'where.exe vim' : 'command -v vim';
-    expect(execSync).toHaveBeenCalledWith(expectedCommand, {
-      stdio: 'ignore',
-    });
+        it(`should return false if "${win32Command}" command does not exist on windows`, () => {
+          Object.defineProperty(process, 'platform', { value: 'win32' });
+          (execSync as Mock).mockImplementation(() => {
+            throw new Error();
+          });
+          expect(checkHasEditorType(editor)).toBe(false);
+        });
+      });
+    }
   });
 
-  it('should return false for vim if "vim" command does not exist', () => {
-    (execSync as Mock).mockImplementation(() => {
-      throw new Error();
+  describe('getDiffCommand', () => {
+    it('should return the correct command for vscode', () => {
+      const command = getDiffCommand('old.txt', 'new.txt', 'vscode');
+      expect(command).toEqual({
+        command: 'code',
+        args: ['--wait', '--diff', 'old.txt', 'new.txt'],
+      });
     });
-    expect(checkHasEditorType('vim')).toBe(false);
-  });
-});
 
-describe('getDiffCommand', () => {
-  it('should return the correct command for vscode', () => {
-    const command = getDiffCommand('old.txt', 'new.txt', 'vscode');
-    expect(command).toEqual({
-      command: 'code',
-      args: ['--wait', '--diff', 'old.txt', 'new.txt'],
+    it('should return the correct command for windsurf', () => {
+      const command = getDiffCommand('old.txt', 'new.txt', 'windsurf');
+      expect(command).toEqual({
+        command: 'windsurf',
+        args: ['--wait', '--diff', 'old.txt', 'new.txt'],
+      });
     });
-  });
-
-  it('should return the correct command for vim', () => {
-    const command = getDiffCommand('old.txt', 'new.txt', 'vim');
-    expect(command?.command).toBe('vim');
-    expect(command?.args).toContain('old.txt');
-    expect(command?.args).toContain('new.txt');
-  });
-
-  it('should return null for an unsupported editor', () => {
-    // @ts-expect-error Testing unsupported editor
-    const command = getDiffCommand('old.txt', 'new.txt', 'foobar');
-    expect(command).toBeNull();
-  });
-});
 
-describe('openDiff', () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-  });
+    it('should return the correct command for cursor', () => {
+      const command = getDiffCommand('old.txt', 'new.txt', 'cursor');
+      expect(command).toEqual({
+        command: 'cursor',
+        args: ['--wait', '--diff', 'old.txt', 'new.txt'],
+      });
+    });
 
-  it('should call spawn for vscode', async () => {
-    const mockSpawn = {
-      on: vi.fn((event, cb) => {
-        if (event === 'close') {
-          cb(0);
-        }
-      }),
-    };
-    (spawn as Mock).mockReturnValue(mockSpawn);
-    await openDiff('old.txt', 'new.txt', 'vscode');
-    expect(spawn).toHaveBeenCalledWith(
-      'code',
-      ['--wait', '--diff', 'old.txt', 'new.txt'],
-      { stdio: 'inherit' },
-    );
-    expect(mockSpawn.on).toHaveBeenCalledWith('close', expect.any(Function));
-    expect(mockSpawn.on).toHaveBeenCalledWith('error', expect.any(Function));
-  });
+    it('should return the correct command for vim', () => {
+      const command = getDiffCommand('old.txt', 'new.txt', 'vim');
+      expect(command).toEqual({
+        command: 'vim',
+        args: [
+          '-d',
+          '-i',
+          'NONE',
+          '-c',
+          'wincmd h | set readonly | wincmd l',
+          '-c',
+          'highlight DiffAdd cterm=bold ctermbg=22 guibg=#005f00 | highlight DiffChange cterm=bold ctermbg=24 guibg=#005f87 | highlight DiffText ctermbg=21 guibg=#0000af | highlight DiffDelete ctermbg=52 guibg=#5f0000',
+          '-c',
+          'set showtabline=2 | set tabline=[Instructions]\\ :wqa(save\\ &\\ quit)\\ \\|\\ i/esc(toggle\\ edit\\ mode)',
+          '-c',
+          'wincmd h | setlocal statusline=OLD\\ FILE',
+          '-c',
+          'wincmd l | setlocal statusline=%#StatusBold#NEW\\ FILE\\ :wqa(save\\ &\\ quit)\\ \\|\\ i/esc(toggle\\ edit\\ mode)',
+          '-c',
+          'autocmd WinClosed * wqa',
+          'old.txt',
+          'new.txt',
+        ],
+      });
+    });
 
-  it('should call execSync for vim', async () => {
-    await openDiff('old.txt', 'new.txt', 'vim');
-    expect(execSync).toHaveBeenCalled();
-    const command = (execSync as Mock).mock.calls[0][0];
-    expect(command).toContain('vim');
-    expect(command).toContain('old.txt');
-    expect(command).toContain('new.txt');
+    it('should return null for an unsupported editor', () => {
+      // @ts-expect-error Testing unsupported editor
+      const command = getDiffCommand('old.txt', 'new.txt', 'foobar');
+      expect(command).toBeNull();
+    });
   });
 
-  it('should handle spawn error for vscode', async () => {
-    const mockSpawn = {
-      on: vi.fn((event, cb) => {
-        if (event === 'error') {
-          cb(new Error('spawn error'));
-        }
-      }),
-    };
-    (spawn as Mock).mockReturnValue(mockSpawn);
-    await expect(openDiff('old.txt', 'new.txt', 'vscode')).rejects.toThrow(
-      'spawn error',
-    );
-  });
-});
+  describe('openDiff', () => {
+    it('should call spawn for vscode', async () => {
+      const mockSpawn = {
+        on: vi.fn((event, cb) => {
+          if (event === 'close') {
+            cb(0);
+          }
+        }),
+      };
+      (spawn as Mock).mockReturnValue(mockSpawn);
+      await openDiff('old.txt', 'new.txt', 'vscode');
+      expect(spawn).toHaveBeenCalledWith(
+        'code',
+        ['--wait', '--diff', 'old.txt', 'new.txt'],
+        { stdio: 'inherit' },
+      );
+      expect(mockSpawn.on).toHaveBeenCalledWith('close', expect.any(Function));
+      expect(mockSpawn.on).toHaveBeenCalledWith('error', expect.any(Function));
+    });
 
-describe('allowEditorTypeInSandbox', () => {
-  afterEach(() => {
-    delete process.env.SANDBOX;
-  });
+    it('should reject if spawn for vscode fails', async () => {
+      const mockError = new Error('spawn error');
+      const mockSpawn = {
+        on: vi.fn((event, cb) => {
+          if (event === 'error') {
+            cb(mockError);
+          }
+        }),
+      };
+      (spawn as Mock).mockReturnValue(mockSpawn);
+      await expect(openDiff('old.txt', 'new.txt', 'vscode')).rejects.toThrow(
+        'spawn error',
+      );
+    });
 
-  it('should allow vim in sandbox mode', () => {
-    process.env.SANDBOX = 'sandbox';
-    expect(allowEditorTypeInSandbox('vim')).toBe(true);
-  });
+    it('should reject if vscode exits with non-zero code', async () => {
+      const mockSpawn = {
+        on: vi.fn((event, cb) => {
+          if (event === 'close') {
+            cb(1);
+          }
+        }),
+      };
+      (spawn as Mock).mockReturnValue(mockSpawn);
+      await expect(openDiff('old.txt', 'new.txt', 'vscode')).rejects.toThrow(
+        'VS Code exited with code 1',
+      );
+    });
 
-  it('should allow vim when not in sandbox mode', () => {
-    delete process.env.SANDBOX;
-    expect(allowEditorTypeInSandbox('vim')).toBe(true);
-  });
+    const execSyncEditors: EditorType[] = ['vim', 'windsurf', 'cursor'];
+    for (const editor of execSyncEditors) {
+      it(`should call execSync for ${editor} on non-windows`, async () => {
+        Object.defineProperty(process, 'platform', { value: 'linux' });
+        await openDiff('old.txt', 'new.txt', editor);
+        expect(execSync).toHaveBeenCalledTimes(1);
+        const diffCommand = getDiffCommand('old.txt', 'new.txt', editor)!;
+        const expectedCommand = `${
+          diffCommand.command
+        } ${diffCommand.args.map((arg) => `"${arg}"`).join(' ')}`;
+        expect(execSync).toHaveBeenCalledWith(expectedCommand, {
+          stdio: 'inherit',
+          encoding: 'utf8',
+        });
+      });
 
-  it('should not allow vscode in sandbox mode', () => {
-    process.env.SANDBOX = 'sandbox';
-    expect(allowEditorTypeInSandbox('vscode')).toBe(false);
-  });
+      it(`should call execSync for ${editor} on windows`, async () => {
+        Object.defineProperty(process, 'platform', { value: 'win32' });
+        await openDiff('old.txt', 'new.txt', editor);
+        expect(execSync).toHaveBeenCalledTimes(1);
+        const diffCommand = getDiffCommand('old.txt', 'new.txt', editor)!;
+        const expectedCommand = `${diffCommand.command} ${diffCommand.args.join(
+          ' ',
+        )}`;
+        expect(execSync).toHaveBeenCalledWith(expectedCommand, {
+          stdio: 'inherit',
+          encoding: 'utf8',
+        });
+      });
+    }
 
-  it('should allow vscode when not in sandbox mode', () => {
-    delete process.env.SANDBOX;
-    expect(allowEditorTypeInSandbox('vscode')).toBe(true);
+    it('should log an error if diff command is not available', async () => {
+      const consoleErrorSpy = vi
+        .spyOn(console, 'error')
+        .mockImplementation(() => {});
+      // @ts-expect-error Testing unsupported editor
+      await openDiff('old.txt', 'new.txt', 'foobar');
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        'No diff tool available. Install vim or vscode.',
+      );
+    });
   });
 
-  it('should not allow windsurf in sandbox mode', () => {
-    process.env.SANDBOX = 'sandbox';
-    expect(allowEditorTypeInSandbox('windsurf')).toBe(false);
-  });
+  describe('allowEditorTypeInSandbox', () => {
+    it('should allow vim in sandbox mode', () => {
+      process.env.SANDBOX = 'sandbox';
+      expect(allowEditorTypeInSandbox('vim')).toBe(true);
+    });
 
-  it('should allow windsurf when not in sandbox mode', () => {
-    delete process.env.SANDBOX;
-    expect(allowEditorTypeInSandbox('windsurf')).toBe(true);
-  });
+    it('should allow vim when not in sandbox mode', () => {
+      expect(allowEditorTypeInSandbox('vim')).toBe(true);
+    });
 
-  it('should not allow cursor in sandbox mode', () => {
-    process.env.SANDBOX = 'sandbox';
-    expect(allowEditorTypeInSandbox('cursor')).toBe(false);
-  });
+    const guiEditors: EditorType[] = ['vscode', 'windsurf', 'cursor'];
+    for (const editor of guiEditors) {
+      it(`should not allow ${editor} in sandbox mode`, () => {
+        process.env.SANDBOX = 'sandbox';
+        expect(allowEditorTypeInSandbox(editor)).toBe(false);
+      });
 
-  it('should allow cursor when not in sandbox mode', () => {
-    delete process.env.SANDBOX;
-    expect(allowEditorTypeInSandbox('cursor')).toBe(true);
+      it(`should allow ${editor} when not in sandbox mode`, () => {
+        expect(allowEditorTypeInSandbox(editor)).toBe(true);
+      });
+    }
   });
-});
 
-describe('isEditorAvailable', () => {
-  afterEach(() => {
-    delete process.env.SANDBOX;
-  });
+  describe('isEditorAvailable', () => {
+    it('should return false for undefined editor', () => {
+      expect(isEditorAvailable(undefined)).toBe(false);
+    });
 
-  it('should return false for undefined editor', () => {
-    expect(isEditorAvailable(undefined)).toBe(false);
-  });
+    it('should return false for empty string editor', () => {
+      expect(isEditorAvailable('')).toBe(false);
+    });
 
-  it('should return false for empty string editor', () => {
-    expect(isEditorAvailable('')).toBe(false);
-  });
+    it('should return false for invalid editor type', () => {
+      expect(isEditorAvailable('invalid-editor')).toBe(false);
+    });
 
-  it('should return false for invalid editor type', () => {
-    expect(isEditorAvailable('invalid-editor')).toBe(false);
-  });
+    it('should return true for vscode when installed and not in sandbox mode', () => {
+      (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/code'));
+      expect(isEditorAvailable('vscode')).toBe(true);
+    });
 
-  it('should return true for vscode when installed and not in sandbox mode', () => {
-    (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/code'));
-    expect(isEditorAvailable('vscode')).toBe(true);
-  });
+    it('should return false for vscode when not installed and not in sandbox mode', () => {
+      (execSync as Mock).mockImplementation(() => {
+        throw new Error();
+      });
+      expect(isEditorAvailable('vscode')).toBe(false);
+    });
 
-  it('should return false for vscode when not installed and not in sandbox mode', () => {
-    (execSync as Mock).mockImplementation(() => {
-      throw new Error();
+    it('should return false for vscode when installed and in sandbox mode', () => {
+      (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/code'));
+      process.env.SANDBOX = 'sandbox';
+      expect(isEditorAvailable('vscode')).toBe(false);
     });
-    expect(isEditorAvailable('vscode')).toBe(false);
-  });
 
-  it('should return false for vscode when installed and in sandbox mode', () => {
-    (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/code'));
-    process.env.SANDBOX = 'sandbox';
-    expect(isEditorAvailable('vscode')).toBe(false);
+    it('should return true for vim when installed and in sandbox mode', () => {
+      (execSync as Mock).mockReturnValue(Buffer.from('/usr/bin/vim'));
+      process.env.SANDBOX = 'sandbox';
+      expect(isEditorAvailable('vim')).toBe(true);
+    });
   });
 });