From 6505b0c8e16f876ba7ca80b19aa3893dc858ce59 Mon Sep 17 00:00:00 2001 From: Arya Gummadi Date: Tue, 19 Aug 2025 17:06:25 -0700 Subject: fix: allow re-auth with another google account (#6544) --- packages/core/src/code_assist/oauth2.test.ts | 79 +++++++++++++++++++++++++++- 1 file changed, 78 insertions(+), 1 deletion(-) (limited to 'packages/core/src/code_assist/oauth2.test.ts') diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts index 25718cb1..4334bd36 100644 --- a/packages/core/src/code_assist/oauth2.test.ts +++ b/packages/core/src/code_assist/oauth2.test.ts @@ -5,7 +5,12 @@ */ import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest'; -import { getOauthClient, resetOauthClientForTesting } from './oauth2.js'; +import { + getOauthClient, + resetOauthClientForTesting, + clearCachedCredentialFile, + clearOauthClientCache, +} from './oauth2.js'; import { getCachedGoogleAccount } from '../utils/user_account.js'; import { OAuth2Client, Compute } from 'google-auth-library'; import * as fs from 'fs'; @@ -510,4 +515,76 @@ describe('oauth2', () => { expect(mockSetCredentials).toHaveBeenCalledWith(cachedCreds); }); }); + + describe('clearCachedCredentialFile', () => { + it('should clear cached credentials and Google account', async () => { + const cachedCreds = { refresh_token: 'test-token' }; + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile(credsPath, JSON.stringify(cachedCreds)); + + const googleAccountPath = path.join( + tempHomeDir, + '.gemini', + 'google_accounts.json', + ); + const accountData = { active: 'test@example.com', old: [] }; + await fs.promises.writeFile( + googleAccountPath, + JSON.stringify(accountData), + ); + + expect(fs.existsSync(credsPath)).toBe(true); + expect(fs.existsSync(googleAccountPath)).toBe(true); + expect(getCachedGoogleAccount()).toBe('test@example.com'); + + await clearCachedCredentialFile(); + expect(fs.existsSync(credsPath)).toBe(false); + expect(getCachedGoogleAccount()).toBeNull(); + const updatedAccountData = JSON.parse( + fs.readFileSync(googleAccountPath, 'utf-8'), + ); + expect(updatedAccountData.active).toBeNull(); + expect(updatedAccountData.old).toContain('test@example.com'); + }); + + it('should clear the in-memory OAuth client cache', async () => { + const mockSetCredentials = vi.fn(); + const mockGetAccessToken = vi + .fn() + .mockResolvedValue({ token: 'test-token' }); + const mockGetTokenInfo = vi.fn().mockResolvedValue({}); + const mockOAuth2Client = { + setCredentials: mockSetCredentials, + getAccessToken: mockGetAccessToken, + getTokenInfo: mockGetTokenInfo, + on: vi.fn(), + } as unknown as OAuth2Client; + (OAuth2Client as unknown as Mock).mockImplementation( + () => mockOAuth2Client, + ); + + // Pre-populate credentials to make getOauthClient resolve quickly + const credsPath = path.join(tempHomeDir, '.gemini', 'oauth_creds.json'); + await fs.promises.mkdir(path.dirname(credsPath), { recursive: true }); + await fs.promises.writeFile( + credsPath, + JSON.stringify({ refresh_token: 'token' }), + ); + + // First call, should create a client + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + expect(OAuth2Client).toHaveBeenCalledTimes(1); + + // Second call, should use cached client + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + expect(OAuth2Client).toHaveBeenCalledTimes(1); + + clearOauthClientCache(); + + // Third call, after clearing cache, should create a new client + await getOauthClient(AuthType.LOGIN_WITH_GOOGLE, mockConfig); + expect(OAuth2Client).toHaveBeenCalledTimes(2); + }); + }); }); -- cgit v1.2.3