/** * @license * Copyright 2025 Google LLC * SPDX-License-Identifier: Apache-2.0 */ import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest'; import { getOauthClient, getCachedGoogleAccountId } from './oauth2.js'; import { OAuth2Client } from 'google-auth-library'; import * as fs from 'fs'; import * as path from 'path'; import http from 'http'; import open from 'open'; import crypto from 'crypto'; import * as os from 'os'; vi.mock('os', async (importOriginal) => { const os = await importOriginal(); return { ...os, homedir: vi.fn(), }; }); vi.mock('google-auth-library'); vi.mock('http'); vi.mock('open'); vi.mock('crypto'); // Mock fetch globally global.fetch = vi.fn(); describe('oauth2', () => { let tempHomeDir: string; beforeEach(() => { tempHomeDir = fs.mkdtempSync( path.join(os.tmpdir(), 'gemini-cli-test-home-'), ); vi.mocked(os.homedir).mockReturnValue(tempHomeDir); }); afterEach(() => { fs.rmSync(tempHomeDir, { recursive: true, force: true }); }); 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 mockGetAccessToken = vi .fn() .mockResolvedValue({ token: 'mock-access-token' }); const mockRefreshAccessToken = vi.fn().mockImplementation((callback) => { // Mock the callback-style refreshAccessToken method const mockTokensWithIdToken = { access_token: 'test-access-token', refresh_token: 'test-refresh-token', id_token: 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiJ0ZXN0LWdvb2dsZS1hY2NvdW50LWlkLTEyMyJ9.signature', // Mock JWT with sub: test-google-account-id-123 }; callback(null, mockTokensWithIdToken); }); const mockVerifyIdToken = vi.fn().mockResolvedValue({ getPayload: () => ({ sub: 'test-google-account-id-123', aud: 'test-audience', iss: 'https://accounts.google.com', }), }); const mockOAuth2Client = { generateAuthUrl: mockGenerateAuthUrl, getToken: mockGetToken, setCredentials: mockSetCredentials, getAccessToken: mockGetAccessToken, refreshAccessToken: mockRefreshAccessToken, verifyIdToken: mockVerifyIdToken, credentials: mockTokens, on: vi.fn(), } as unknown as OAuth2Client; vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client); vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never); vi.mocked(open).mockImplementation(async () => ({}) as never); // Mock the UserInfo API response vi.mocked(global.fetch).mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue({ id: 'test-google-account-id-123' }), } as unknown as Response); let requestCallback!: http.RequestListener< typeof http.IncomingMessage, typeof http.ServerResponse >; let serverListeningCallback: (value: unknown) => void; const serverListeningPromise = new Promise( (resolve) => (serverListeningCallback = resolve), ); let capturedPort = 0; const mockHttpServer = { listen: vi.fn((port: number, callback?: () => void) => { capturedPort = port; if (callback) { callback(); } serverListeningCallback(undefined); }), close: vi.fn((callback?: () => void) => { if (callback) { callback(); } }), on: vi.fn(), address: () => ({ port: capturedPort }), }; vi.mocked(http.createServer).mockImplementation((cb) => { requestCallback = cb as http.RequestListener< typeof http.IncomingMessage, typeof http.ServerResponse >; return mockHttpServer as unknown as http.Server; }); const clientPromise = getOauthClient(); // wait for server to start listening. await serverListeningPromise; const mockReq = { url: `/oauth2callback?code=${mockCode}&state=${mockState}`, } as http.IncomingMessage; const mockRes = { writeHead: vi.fn(), end: vi.fn(), } as unknown as http.ServerResponse; await requestCallback(mockReq, mockRes); const client = await clientPromise; expect(client).toBe(mockOAuth2Client); expect(open).toHaveBeenCalledWith(mockAuthUrl); expect(mockGetToken).toHaveBeenCalledWith({ code: mockCode, redirect_uri: `http://localhost:${capturedPort}/oauth2callback`, }); expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens); // Verify Google Account ID was cached const googleAccountIdPath = path.join( tempHomeDir, '.gemini', 'google_account_id', ); expect(fs.existsSync(googleAccountIdPath)).toBe(true); const cachedGoogleAccountId = fs.readFileSync(googleAccountIdPath, 'utf-8'); expect(cachedGoogleAccountId).toBe('test-google-account-id-123'); // Verify the getCachedGoogleAccountId function works expect(getCachedGoogleAccountId()).toBe('test-google-account-id-123'); }); });