From 197704c630cbf85aa4d113c7a62e305827043910 Mon Sep 17 00:00:00 2001 From: "N. Taylor Mullen" Date: Sun, 15 Jun 2025 22:41:32 -0700 Subject: feat(test): Increase test coverage across CLI and Core packages (#1089) --- packages/core/src/code_assist/oauth2.test.ts | 92 ++++++++++++++++ packages/core/src/code_assist/server.test.ts | 151 +++++++++++++++++++++++++++ 2 files changed, 243 insertions(+) create mode 100644 packages/core/src/code_assist/oauth2.test.ts create mode 100644 packages/core/src/code_assist/server.test.ts (limited to 'packages/core/src/code_assist') 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(); + }); +}); -- cgit v1.2.3