From 21fba832d1b4ea7af43fb887d9b2b38fcf8210d0 Mon Sep 17 00:00:00 2001 From: Tommaso Sciortino Date: Fri, 30 May 2025 18:25:47 -0700 Subject: Rename server->core (#638) --- packages/core/src/utils/nextSpeakerChecker.test.ts | 235 +++++++++++++++++++++ 1 file changed, 235 insertions(+) create mode 100644 packages/core/src/utils/nextSpeakerChecker.test.ts (limited to 'packages/core/src/utils/nextSpeakerChecker.test.ts') diff --git a/packages/core/src/utils/nextSpeakerChecker.test.ts b/packages/core/src/utils/nextSpeakerChecker.test.ts new file mode 100644 index 00000000..872e00f6 --- /dev/null +++ b/packages/core/src/utils/nextSpeakerChecker.test.ts @@ -0,0 +1,235 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, vi, beforeEach, Mock, afterEach } from 'vitest'; +import { Content, GoogleGenAI, Models } from '@google/genai'; +import { GeminiClient } from '../core/client.js'; +import { Config } from '../config/config.js'; +import { checkNextSpeaker, NextSpeakerResponse } from './nextSpeakerChecker.js'; +import { GeminiChat } from '../core/geminiChat.js'; + +// Mock GeminiClient and Config constructor +vi.mock('../core/client.js'); +vi.mock('../config/config.js'); + +// Define mocks for GoogleGenAI and Models instances that will be used across tests +const mockModelsInstance = { + generateContent: vi.fn(), + generateContentStream: vi.fn(), + countTokens: vi.fn(), + embedContent: vi.fn(), + batchEmbedContents: vi.fn(), +} as unknown as Models; + +const mockGoogleGenAIInstance = { + getGenerativeModel: vi.fn().mockReturnValue(mockModelsInstance), + // Add other methods of GoogleGenAI if they are directly used by GeminiChat constructor or its methods +} as unknown as GoogleGenAI; + +vi.mock('@google/genai', async () => { + const actualGenAI = + await vi.importActual('@google/genai'); + return { + ...actualGenAI, + GoogleGenAI: vi.fn(() => mockGoogleGenAIInstance), // Mock constructor to return the predefined instance + // If Models is instantiated directly in GeminiChat, mock its constructor too + // For now, assuming Models instance is obtained via getGenerativeModel + }; +}); + +describe('checkNextSpeaker', () => { + let chatInstance: GeminiChat; + let mockGeminiClient: GeminiClient; + let MockConfig: Mock; + const abortSignal = new AbortController().signal; + + beforeEach(() => { + MockConfig = vi.mocked(Config); + const mockConfigInstance = new MockConfig( + 'test-api-key', + 'gemini-pro', + false, + '.', + false, + undefined, + false, + undefined, + undefined, + undefined, + ); + + mockGeminiClient = new GeminiClient(mockConfigInstance); + + // Reset mocks before each test to ensure test isolation + vi.mocked(mockModelsInstance.generateContent).mockReset(); + vi.mocked(mockModelsInstance.generateContentStream).mockReset(); + + // GeminiChat will receive the mocked instances via the mocked GoogleGenAI constructor + chatInstance = new GeminiChat( + mockGoogleGenAIInstance, // This will be the instance returned by the mocked GoogleGenAI constructor + mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel + 'gemini-pro', // model name + {}, + [], // initial history + ); + + // Spy on getHistory for chatInstance + vi.spyOn(chatInstance, 'getHistory'); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it('should return null if history is empty', async () => { + (chatInstance.getHistory as Mock).mockReturnValue([]); + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toBeNull(); + expect(mockGeminiClient.generateJson).not.toHaveBeenCalled(); + }); + + it('should return null if the last speaker was the user', async () => { + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'user', parts: [{ text: 'Hello' }] }, + ] as Content[]); + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toBeNull(); + expect(mockGeminiClient.generateJson).not.toHaveBeenCalled(); + }); + + it("should return { next_speaker: 'model' } when model intends to continue", async () => { + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'model', parts: [{ text: 'I will now do something.' }] }, + ] as Content[]); + const mockApiResponse: NextSpeakerResponse = { + reasoning: 'Model stated it will do something.', + next_speaker: 'model', + }; + (mockGeminiClient.generateJson as Mock).mockResolvedValue(mockApiResponse); + + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toEqual(mockApiResponse); + expect(mockGeminiClient.generateJson).toHaveBeenCalledTimes(1); + }); + + it("should return { next_speaker: 'user' } when model asks a question", async () => { + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'model', parts: [{ text: 'What would you like to do?' }] }, + ] as Content[]); + const mockApiResponse: NextSpeakerResponse = { + reasoning: 'Model asked a question.', + next_speaker: 'user', + }; + (mockGeminiClient.generateJson as Mock).mockResolvedValue(mockApiResponse); + + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toEqual(mockApiResponse); + }); + + it("should return { next_speaker: 'user' } when model makes a statement", async () => { + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'model', parts: [{ text: 'This is a statement.' }] }, + ] as Content[]); + const mockApiResponse: NextSpeakerResponse = { + reasoning: 'Model made a statement, awaiting user input.', + next_speaker: 'user', + }; + (mockGeminiClient.generateJson as Mock).mockResolvedValue(mockApiResponse); + + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toEqual(mockApiResponse); + }); + + it('should return null if geminiClient.generateJson throws an error', async () => { + const consoleWarnSpy = vi + .spyOn(console, 'warn') + .mockImplementation(() => {}); + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'model', parts: [{ text: 'Some model output.' }] }, + ] as Content[]); + (mockGeminiClient.generateJson as Mock).mockRejectedValue( + new Error('API Error'), + ); + + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toBeNull(); + consoleWarnSpy.mockRestore(); + }); + + it('should return null if geminiClient.generateJson returns invalid JSON (missing next_speaker)', async () => { + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'model', parts: [{ text: 'Some model output.' }] }, + ] as Content[]); + (mockGeminiClient.generateJson as Mock).mockResolvedValue({ + reasoning: 'This is incomplete.', + } as unknown as NextSpeakerResponse); // Type assertion to simulate invalid response + + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toBeNull(); + }); + + it('should return null if geminiClient.generateJson returns a non-string next_speaker', async () => { + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'model', parts: [{ text: 'Some model output.' }] }, + ] as Content[]); + (mockGeminiClient.generateJson as Mock).mockResolvedValue({ + reasoning: 'Model made a statement, awaiting user input.', + next_speaker: 123, // Invalid type + } as unknown as NextSpeakerResponse); + + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toBeNull(); + }); + + it('should return null if geminiClient.generateJson returns an invalid next_speaker string value', async () => { + (chatInstance.getHistory as Mock).mockReturnValue([ + { role: 'model', parts: [{ text: 'Some model output.' }] }, + ] as Content[]); + (mockGeminiClient.generateJson as Mock).mockResolvedValue({ + reasoning: 'Model made a statement, awaiting user input.', + next_speaker: 'neither', // Invalid enum value + } as unknown as NextSpeakerResponse); + + const result = await checkNextSpeaker( + chatInstance, + mockGeminiClient, + abortSignal, + ); + expect(result).toBeNull(); + }); +}); -- cgit v1.2.3