fix: add privacy settings hook and tests (#6360)

This commit is contained in:
Arya Gummadi 2025-08-18 23:57:10 -07:00 committed by GitHub
parent 8f8082fe3d
commit ec0d9f4ff7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 261 additions and 11 deletions

View File

@ -0,0 +1,242 @@
/**
* @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<string, unknown>,
_sessionId?: string,
_userTier?: UserTierId,
) {}
}
class MockLoggingContentGenerator {
getWrapped = vi.fn();
constructor(
_wrapped?: ContentGenerator,
_config?: Record<string, unknown>,
) {}
}
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<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);
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<typeof vi.fn>
).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<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);
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,
});
});
});

View File

@ -5,7 +5,12 @@
*/
import { useState, useEffect, useCallback } from 'react';
import { Config, CodeAssistServer, UserTierId } from '@google/gemini-cli-core';
import {
Config,
CodeAssistServer,
UserTierId,
LoggingContentGenerator,
} from '@google/gemini-cli-core';
export interface PrivacyState {
isLoading: boolean;
@ -80,7 +85,13 @@ export const usePrivacySettings = (config: Config) => {
};
function getCodeAssistServer(config: Config): CodeAssistServer {
const server = config.getGeminiClient().getContentGenerator();
let server = config.getGeminiClient().getContentGenerator();
// Unwrap LoggingContentGenerator if present
if (server instanceof LoggingContentGenerator) {
server = server.getWrapped();
}
// Neither of these cases should ever happen.
if (!(server instanceof CodeAssistServer)) {
throw new Error('Oauth not being used');

View File

@ -27,20 +27,12 @@ import {
} from '../telemetry/loggers.js';
import { ContentGenerator } from './contentGenerator.js';
import { toContents } from '../code_assist/converter.js';
import { isStructuredError } from '../utils/quotaErrorDetection.js';
interface StructuredError {
status: number;
}
export function isStructuredError(error: unknown): error is StructuredError {
return (
typeof error === 'object' &&
error !== null &&
'status' in error &&
typeof (error as StructuredError).status === 'number'
);
}
/**
* A decorator that wraps a ContentGenerator to add logging to API calls.
*/
@ -50,6 +42,10 @@ export class LoggingContentGenerator implements ContentGenerator {
private readonly config: Config,
) {}
getWrapped(): ContentGenerator {
return this.wrapped;
}
private logApiRequest(
contents: Content[],
model: string,

View File

@ -10,6 +10,7 @@ export * from './config/config.js';
// Export Core Logic
export * from './core/client.js';
export * from './core/contentGenerator.js';
export * from './core/loggingContentGenerator.js';
export * from './core/geminiChat.js';
export * from './core/logger.js';
export * from './core/prompts.js';