Enable citations by default for certain users. (#7438)

packages/cli/src/ui/hooks/usePrivacySettings.test.ts

@@ -8,150 +8,31 @@ import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { renderHook, waitFor } from '@testing-library/react';
 import type {
   Config,
-  GeminiClient,
-  ContentGenerator,
-} from '@google/gemini-cli-core';
-import {
   CodeAssistServer,
-  LoggingContentGenerator,
-  UserTierId,
+  LoadCodeAssistResponse,
 } from '@google/gemini-cli-core';
-import type { OAuth2Client } from 'google-auth-library';
+import { UserTierId, getCodeAssistServer } from '@google/gemini-cli-core';
 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<string, unknown>,
-      _sessionId?: string,
-      _userTier?: UserTierId,
-    ) {}
-  }
-
-  class MockLoggingContentGenerator {
-    getWrapped = vi.fn();
-
-    constructor(
-      _wrapped?: ContentGenerator,
-      _config?: Record<string, unknown>,
-    ) {}
-  }
-
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import('@google/gemini-cli-core')>();
   return {
-    Config: vi.fn(),
-    CodeAssistServer: MockCodeAssistServer,
-    LoggingContentGenerator: MockLoggingContentGenerator,
-    GeminiClient: vi.fn(),
-    UserTierId: {
-      FREE: 'free-tier',
-      LEGACY: 'legacy-tier',
-      STANDARD: 'standard-tier',
-    },
+    ...actual,
+    getCodeAssistServer: vi.fn(),
   };
 });
 
 describe('usePrivacySettings', () => {
-  let mockConfig: Config;
-  let mockClient: GeminiClient;
-  let mockCodeAssistServer: CodeAssistServer;
-  let mockLoggingContentGenerator: LoggingContentGenerator;
+  const mockConfig = {} as unknown as Config;
 
   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<typeof vi.fn>
-    ).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<typeof vi.fn>
-    ).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<typeof vi.fn>).mockResolvedValue(
-      {
-        currentTier: { id: UserTierId.FREE },
-      },
-    );
-    (
-      directServer.getCodeAssistGlobalUserSetting as ReturnType<typeof vi.fn>
-    ).mockResolvedValue({
-      freeTierDataCollectionOptin: true,
-    });
-
-    mockClient.getContentGenerator = vi.fn().mockReturnValue(directServer);
+  it('should throw error when content generator is not a CodeAssistServer', async () => {
+    vi.mocked(getCodeAssistServer).mockReturnValue(undefined);
 
     const { result } = renderHook(() => usePrivacySettings(mockConfig));
 
@@ -159,18 +40,18 @@ describe('usePrivacySettings', () => {
       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);
+    expect(result.current.privacyState.error).toBe('Oauth not being used');
   });
 
   it('should handle paid tier users correctly', async () => {
     // Mock paid tier response
-    (
-      mockCodeAssistServer.loadCodeAssist as ReturnType<typeof vi.fn>
-    ).mockResolvedValue({
-      currentTier: { id: UserTierId.STANDARD },
-    });
+    vi.mocked(getCodeAssistServer).mockReturnValue({
+      projectId: 'test-project-id',
+      loadCodeAssist: () =>
+        ({
+          currentTier: { id: UserTierId.STANDARD },
+        }) as unknown as LoadCodeAssistResponse,
+    } as unknown as CodeAssistServer);
 
     const { result } = renderHook(() => usePrivacySettings(mockConfig));
 
@@ -183,31 +64,13 @@ describe('usePrivacySettings', () => {
     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<typeof vi.fn>
-    ).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<typeof vi.fn>
-    ).mockReturnValue(mockServerNoProject);
+    vi.mocked(getCodeAssistServer).mockReturnValue({
+      loadCodeAssist: () =>
+        ({
+          currentTier: { id: UserTierId.FREE },
+        }) as unknown as LoadCodeAssistResponse,
+    } as unknown as CodeAssistServer);
 
     const { result } = renderHook(() => usePrivacySettings(mockConfig));
 
@@ -215,10 +78,27 @@ describe('usePrivacySettings', () => {
       expect(result.current.privacyState.isLoading).toBe(false);
     });
 
-    expect(result.current.privacyState.error).toBe('Oauth not being used');
+    expect(result.current.privacyState.error).toBe(
+      'CodeAssist server is missing a project ID',
+    );
   });
 
   it('should update data collection opt-in setting', async () => {
+    const mockCodeAssistServer = {
+      projectId: 'test-project-id',
+      getCodeAssistGlobalUserSetting: vi.fn().mockResolvedValue({
+        freeTierDataCollectionOptin: true,
+      }),
+      setCodeAssistGlobalUserSetting: vi.fn().mockResolvedValue({
+        freeTierDataCollectionOptin: false,
+      }),
+      loadCodeAssist: () =>
+        ({
+          currentTier: { id: UserTierId.FREE },
+        }) as unknown as LoadCodeAssistResponse,
+    } as unknown as CodeAssistServer;
+    vi.mocked(getCodeAssistServer).mockReturnValue(mockCodeAssistServer);
+
     const { result } = renderHook(() => usePrivacySettings(mockConfig));
 
     // Wait for initial load