/** * @license * Copyright 2025 Google LLC * SPDX-License-Identifier: Apache-2.0 */ import { describe, it, expect, beforeEach, vi } from 'vitest'; import { renderHook, waitFor } from '@testing-library/react'; import { Config, CodeAssistServer, LoggingContentGenerator, UserTierId, GeminiClient, ContentGenerator, } from '@google/gemini-cli-core'; import { OAuth2Client } from 'google-auth-library'; import { usePrivacySettings } from './usePrivacySettings.js'; // Mock the dependencies vi.mock('@google/gemini-cli-core', () => { // Mock classes for instanceof checks class MockCodeAssistServer { projectId = 'test-project-id'; loadCodeAssist = vi.fn(); getCodeAssistGlobalUserSetting = vi.fn(); setCodeAssistGlobalUserSetting = vi.fn(); constructor( _client?: GeminiClient, _projectId?: string, _httpOptions?: Record, _sessionId?: string, _userTier?: UserTierId, ) {} } class MockLoggingContentGenerator { getWrapped = vi.fn(); constructor( _wrapped?: ContentGenerator, _config?: Record, ) {} } return { Config: vi.fn(), CodeAssistServer: MockCodeAssistServer, LoggingContentGenerator: MockLoggingContentGenerator, GeminiClient: vi.fn(), UserTierId: { FREE: 'free-tier', LEGACY: 'legacy-tier', STANDARD: 'standard-tier', }, }; }); describe('usePrivacySettings', () => { let mockConfig: Config; let mockClient: GeminiClient; let mockCodeAssistServer: CodeAssistServer; let mockLoggingContentGenerator: LoggingContentGenerator; beforeEach(() => { vi.clearAllMocks(); // Create mock CodeAssistServer instance mockCodeAssistServer = new CodeAssistServer( null as unknown as OAuth2Client, 'test-project-id', ) as unknown as CodeAssistServer; ( mockCodeAssistServer.loadCodeAssist as ReturnType ).mockResolvedValue({ currentTier: { id: UserTierId.FREE }, }); ( mockCodeAssistServer.getCodeAssistGlobalUserSetting as ReturnType< typeof vi.fn > ).mockResolvedValue({ freeTierDataCollectionOptin: true, }); ( mockCodeAssistServer.setCodeAssistGlobalUserSetting as ReturnType< typeof vi.fn > ).mockResolvedValue({ freeTierDataCollectionOptin: false, }); // Create mock LoggingContentGenerator that wraps the CodeAssistServer mockLoggingContentGenerator = new LoggingContentGenerator( mockCodeAssistServer, null as unknown as Config, ) as unknown as LoggingContentGenerator; ( mockLoggingContentGenerator.getWrapped as ReturnType ).mockReturnValue(mockCodeAssistServer); // Create mock GeminiClient mockClient = { getContentGenerator: vi.fn().mockReturnValue(mockLoggingContentGenerator), } as unknown as GeminiClient; // Create mock Config mockConfig = { getGeminiClient: vi.fn().mockReturnValue(mockClient), } as unknown as Config; }); it('should handle LoggingContentGenerator wrapper correctly and not throw "Oauth not being used" error', async () => { const { result } = renderHook(() => usePrivacySettings(mockConfig)); // Initial state should be loading expect(result.current.privacyState.isLoading).toBe(true); expect(result.current.privacyState.error).toBeUndefined(); // Wait for the hook to complete await waitFor(() => { expect(result.current.privacyState.isLoading).toBe(false); }); // Should not have the "Oauth not being used" error expect(result.current.privacyState.error).toBeUndefined(); expect(result.current.privacyState.isFreeTier).toBe(true); expect(result.current.privacyState.dataCollectionOptIn).toBe(true); // Verify that getWrapped was called to unwrap the LoggingContentGenerator expect(mockLoggingContentGenerator.getWrapped).toHaveBeenCalled(); }); it('should work with direct CodeAssistServer (no wrapper)', async () => { // Test case where the content generator is directly a CodeAssistServer const directServer = new CodeAssistServer( null as unknown as OAuth2Client, 'test-project-id', ) as unknown as CodeAssistServer; (directServer.loadCodeAssist as ReturnType).mockResolvedValue( { currentTier: { id: UserTierId.FREE }, }, ); ( directServer.getCodeAssistGlobalUserSetting as ReturnType ).mockResolvedValue({ freeTierDataCollectionOptin: true, }); mockClient.getContentGenerator = vi.fn().mockReturnValue(directServer); const { result } = renderHook(() => usePrivacySettings(mockConfig)); await waitFor(() => { expect(result.current.privacyState.isLoading).toBe(false); }); expect(result.current.privacyState.error).toBeUndefined(); expect(result.current.privacyState.isFreeTier).toBe(true); expect(result.current.privacyState.dataCollectionOptIn).toBe(true); }); it('should handle paid tier users correctly', async () => { // Mock paid tier response ( mockCodeAssistServer.loadCodeAssist as ReturnType ).mockResolvedValue({ currentTier: { id: UserTierId.STANDARD }, }); const { result } = renderHook(() => usePrivacySettings(mockConfig)); await waitFor(() => { expect(result.current.privacyState.isLoading).toBe(false); }); expect(result.current.privacyState.error).toBeUndefined(); expect(result.current.privacyState.isFreeTier).toBe(false); expect(result.current.privacyState.dataCollectionOptIn).toBeUndefined(); }); it('should throw error when content generator is not a CodeAssistServer', async () => { // Mock a non-CodeAssistServer content generator const mockOtherGenerator = { someOtherMethod: vi.fn() }; ( mockLoggingContentGenerator.getWrapped as ReturnType ).mockReturnValue(mockOtherGenerator); const { result } = renderHook(() => usePrivacySettings(mockConfig)); await waitFor(() => { expect(result.current.privacyState.isLoading).toBe(false); }); expect(result.current.privacyState.error).toBe('Oauth not being used'); }); it('should throw error when CodeAssistServer has no projectId', async () => { // Mock CodeAssistServer without projectId const mockServerNoProject = { ...mockCodeAssistServer, projectId: undefined, }; ( mockLoggingContentGenerator.getWrapped as ReturnType ).mockReturnValue(mockServerNoProject); const { result } = renderHook(() => usePrivacySettings(mockConfig)); await waitFor(() => { expect(result.current.privacyState.isLoading).toBe(false); }); expect(result.current.privacyState.error).toBe('Oauth not being used'); }); it('should update data collection opt-in setting', async () => { const { result } = renderHook(() => usePrivacySettings(mockConfig)); // Wait for initial load await waitFor(() => { expect(result.current.privacyState.isLoading).toBe(false); }); // Update the setting await result.current.updateDataCollectionOptIn(false); // Wait for update to complete await waitFor(() => { expect(result.current.privacyState.dataCollectionOptIn).toBe(false); }); expect( mockCodeAssistServer.setCodeAssistGlobalUserSetting, ).toHaveBeenCalledWith({ cloudaicompanionProject: 'test-project-id', freeTierDataCollectionOptin: false, }); }); });