feat(test): Increase test coverage across CLI and Core packages (#1089)

2 files changed, 243 insertions, 0 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();
+  });
+});