1 files changed, 16 insertions, 29 deletions
diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index d8cd525b..cec05121 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -5,7 +5,8 @@
  */
 
 import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
-import { getOauthClient, getCachedGoogleAccountId } from './oauth2.js';
+import { getOauthClient } from './oauth2.js';
+import { getCachedGoogleAccount } from '../utils/user_account.js';
 import { OAuth2Client, Compute } from 'google-auth-library';
 import * as fs from 'fs';
 import * as path from 'path';
@@ -66,30 +67,11 @@ describe('oauth2', () => {
     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;
@@ -103,7 +85,9 @@ describe('oauth2', () => {
     // Mock the UserInfo API response
     (global.fetch as Mock).mockResolvedValue({
       ok: true,
-      json: vi.fn().mockResolvedValue({ id: 'test-google-account-id-123' }),
+      json: vi
+        .fn()
+        .mockResolvedValue({ email: '[email protected]' }),
     } as unknown as Response);
 
     let requestCallback!: http.RequestListener<
@@ -169,18 +153,21 @@ describe('oauth2', () => {
     });
     expect(mockSetCredentials).toHaveBeenCalledWith(mockTokens);
 
-    // Verify Google Account ID was cached
-    const googleAccountIdPath = path.join(
+    // Verify Google Account was cached
+    const googleAccountPath = path.join(
       tempHomeDir,
       '.gemini',
-      'google_account_id',
+      'google_accounts.json',
     );
-    expect(fs.existsSync(googleAccountIdPath)).toBe(true);
-    const cachedGoogleAccountId = fs.readFileSync(googleAccountIdPath, 'utf-8');
-    expect(cachedGoogleAccountId).toBe('test-google-account-id-123');
+    expect(fs.existsSync(googleAccountPath)).toBe(true);
+    const cachedGoogleAccount = fs.readFileSync(googleAccountPath, 'utf-8');
+    expect(JSON.parse(cachedGoogleAccount)).toEqual({
+      active: '[email protected]',
+      old: [],
+    });
 
-    // Verify the getCachedGoogleAccountId function works
-    expect(getCachedGoogleAccountId()).toBe('test-google-account-id-123');
+    // Verify the getCachedGoogleAccount function works
+    expect(getCachedGoogleAccount()).toBe('[email protected]');
   });
 
   describe('in Cloud Shell', () => {