fix: add privacy settings hook and tests (#6360)
This commit is contained in:
parent
8f8082fe3d
commit
ec0d9f4ff7
|
@ -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,
|
||||
});
|
||||
});
|
||||
});
|
|
@ -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');
|
||||
|
|
|
@ -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,
|
||||
|
|
|
@ -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';
|
||||
|
|
Loading…
Reference in New Issue