1 files changed, 96 insertions, 8 deletions

diff --git a/packages/core/src/code_assist/oauth2.test.ts b/packages/core/src/code_assist/oauth2.test.ts
index ae5d80d6..76d43726 100644
--- a/packages/core/src/code_assist/oauth2.test.ts
+++ b/packages/core/src/code_assist/oauth2.test.ts
@@ -4,15 +4,16 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { describe, it, expect, vi, beforeEach, afterEach, Mock } from 'vitest';
 import { getOauthClient, getCachedGoogleAccountId } from './oauth2.js';
-import { OAuth2Client } from 'google-auth-library';
+import { OAuth2Client, Compute } 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';
+import { AuthType } from '../core/contentGenerator.js';
 
 vi.mock('os', async (importOriginal) => {
   const os = await importOriginal<typeof import('os')>();
@@ -37,10 +38,12 @@ describe('oauth2', () => {
     tempHomeDir = fs.mkdtempSync(
       path.join(os.tmpdir(), 'gemini-cli-test-home-'),
     );
-    vi.mocked(os.homedir).mockReturnValue(tempHomeDir);
+    (os.homedir as Mock).mockReturnValue(tempHomeDir);
   });
   afterEach(() => {
     fs.rmSync(tempHomeDir, { recursive: true, force: true });
+    vi.clearAllMocks();
+    delete process.env.CLOUD_SHELL;
   });
 
   it('should perform a web login', async () => {
@@ -85,13 +88,15 @@ describe('oauth2', () => {
       credentials: mockTokens,
       on: vi.fn(),
     } as unknown as OAuth2Client;
-    vi.mocked(OAuth2Client).mockImplementation(() => mockOAuth2Client);
+    (OAuth2Client as unknown as Mock).mockImplementation(
+      () => mockOAuth2Client,
+    );
 
     vi.spyOn(crypto, 'randomBytes').mockReturnValue(mockState as never);
-    vi.mocked(open).mockImplementation(async () => ({}) as never);
+    (open as Mock).mockImplementation(async () => ({}) as never);
 
     // Mock the UserInfo API response
-    vi.mocked(global.fetch).mockResolvedValue({
+    (global.fetch as Mock).mockResolvedValue({
       ok: true,
       json: vi.fn().mockResolvedValue({ id: 'test-google-account-id-123' }),
     } as unknown as Response);
@@ -123,7 +128,7 @@ describe('oauth2', () => {
       on: vi.fn(),
       address: () => ({ port: capturedPort }),
     };
-    vi.mocked(http.createServer).mockImplementation((cb) => {
+    (http.createServer as Mock).mockImplementation((cb) => {
       requestCallback = cb as http.RequestListener<
         typeof http.IncomingMessage,
         typeof http.ServerResponse
@@ -131,7 +136,7 @@ describe('oauth2', () => {
       return mockHttpServer as unknown as http.Server;
     });
 
-    const clientPromise = getOauthClient();
+    const clientPromise = getOauthClient(AuthType.LOGIN_WITH_GOOGLE);
 
     // wait for server to start listening.
     await serverListeningPromise;
@@ -169,4 +174,87 @@ describe('oauth2', () => {
     // Verify the getCachedGoogleAccountId function works
     expect(getCachedGoogleAccountId()).toBe('test-google-account-id-123');
   });
+
+  describe('in Cloud Shell', () => {
+    const mockGetAccessToken = vi.fn();
+    let mockComputeClient: Compute;
+
+    beforeEach(() => {
+      vi.spyOn(os, 'homedir').mockReturnValue('/user/home');
+      vi.spyOn(fs.promises, 'mkdir').mockResolvedValue(undefined);
+      vi.spyOn(fs.promises, 'writeFile').mockResolvedValue(undefined);
+      vi.spyOn(fs.promises, 'readFile').mockRejectedValue(
+        new Error('File not found'),
+      ); // Default to no cached creds
+
+      mockGetAccessToken.mockResolvedValue({ token: 'test-access-token' });
+      mockComputeClient = {
+        credentials: { refresh_token: 'test-refresh-token' },
+        getAccessToken: mockGetAccessToken,
+      } as unknown as Compute;
+
+      (Compute as unknown as Mock).mockImplementation(() => mockComputeClient);
+    });
+
+    it('should attempt to load cached credentials first', async () => {
+      const cachedCreds = { refresh_token: 'cached-token' };
+      vi.spyOn(fs.promises, 'readFile').mockResolvedValue(
+        JSON.stringify(cachedCreds),
+      );
+
+      const mockClient = {
+        setCredentials: vi.fn(),
+        getAccessToken: vi.fn().mockResolvedValue({ token: 'test-token' }),
+        getTokenInfo: vi.fn().mockResolvedValue({}),
+        on: vi.fn(),
+      };
+
+      // To mock the new OAuth2Client() inside the function
+      (OAuth2Client as unknown as Mock).mockImplementation(
+        () => mockClient as unknown as OAuth2Client,
+      );
+
+      await getOauthClient(AuthType.LOGIN_WITH_GOOGLE);
+
+      expect(fs.promises.readFile).toHaveBeenCalledWith(
+        '/user/home/.gemini/oauth_creds.json',
+        'utf-8',
+      );
+      expect(mockClient.setCredentials).toHaveBeenCalledWith(cachedCreds);
+      expect(mockClient.getAccessToken).toHaveBeenCalled();
+      expect(mockClient.getTokenInfo).toHaveBeenCalled();
+      expect(Compute).not.toHaveBeenCalled(); // Should not fetch new client if cache is valid
+    });
+
+    it('should use Compute to get a client if no cached credentials exist', async () => {
+      await getOauthClient(AuthType.CLOUD_SHELL);
+
+      expect(Compute).toHaveBeenCalledWith({});
+      expect(mockGetAccessToken).toHaveBeenCalled();
+    });
+
+    it('should not cache the credentials after fetching them via ADC', async () => {
+      const newCredentials = { refresh_token: 'new-adc-token' };
+      mockComputeClient.credentials = newCredentials;
+      mockGetAccessToken.mockResolvedValue({ token: 'new-adc-token' });
+
+      await getOauthClient(AuthType.CLOUD_SHELL);
+
+      expect(fs.promises.writeFile).not.toHaveBeenCalled();
+    });
+
+    it('should return the Compute client on successful ADC authentication', async () => {
+      const client = await getOauthClient(AuthType.CLOUD_SHELL);
+      expect(client).toBe(mockComputeClient);
+    });
+
+    it('should throw an error if ADC fails', async () => {
+      const testError = new Error('ADC Failed');
+      mockGetAccessToken.mockRejectedValue(testError);
+
+      await expect(getOauthClient(AuthType.CLOUD_SHELL)).rejects.toThrow(
+        'Could not authenticate using Cloud Shell credentials. Please select a different authentication method or ensure you are in a properly configured environment. Error: ADC Failed',
+      );
+    });
+  });
 });