@quantumai/quantum-cli-core
Version:
Quantum CLI Core - Multi-LLM Collaboration System
152 lines • 7.38 kB
JavaScript
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import { GeminiClient } from '../core/client.js';
import { Config } from '../config/config.js';
import { checkNextSpeaker } 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(),
};
const mockGoogleGenAIInstance = {
getGenerativeModel: vi.fn().mockReturnValue(mockModelsInstance),
// Add other methods of GoogleGenAI if they are directly used by GeminiChat constructor or its methods
};
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;
let mockGeminiClient;
let MockConfig;
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(mockConfigInstance, mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel
{}, []);
// Spy on getHistory for chatInstance
vi.spyOn(chatInstance, 'getHistory');
});
afterEach(() => {
vi.clearAllMocks();
});
it('should return null if history is empty', async () => {
chatInstance.getHistory.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.mockReturnValue([
{ role: 'user', parts: [{ text: 'Hello' }] },
]);
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.mockReturnValue([
{ role: 'model', parts: [{ text: 'I will now do something.' }] },
]);
const mockApiResponse = {
reasoning: 'Model stated it will do something.',
next_speaker: 'model',
};
mockGeminiClient.generateJson.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.mockReturnValue([
{ role: 'model', parts: [{ text: 'What would you like to do?' }] },
]);
const mockApiResponse = {
reasoning: 'Model asked a question.',
next_speaker: 'user',
};
mockGeminiClient.generateJson.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.mockReturnValue([
{ role: 'model', parts: [{ text: 'This is a statement.' }] },
]);
const mockApiResponse = {
reasoning: 'Model made a statement, awaiting user input.',
next_speaker: 'user',
};
mockGeminiClient.generateJson.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.mockReturnValue([
{ role: 'model', parts: [{ text: 'Some model output.' }] },
]);
mockGeminiClient.generateJson.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.mockReturnValue([
{ role: 'model', parts: [{ text: 'Some model output.' }] },
]);
mockGeminiClient.generateJson.mockResolvedValue({
reasoning: 'This is incomplete.',
}); // 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.mockReturnValue([
{ role: 'model', parts: [{ text: 'Some model output.' }] },
]);
mockGeminiClient.generateJson.mockResolvedValue({
reasoning: 'Model made a statement, awaiting user input.',
next_speaker: 123, // Invalid type
});
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.mockReturnValue([
{ role: 'model', parts: [{ text: 'Some model output.' }] },
]);
mockGeminiClient.generateJson.mockResolvedValue({
reasoning: 'Model made a statement, awaiting user input.',
next_speaker: 'neither', // Invalid enum value
});
const result = await checkNextSpeaker(chatInstance, mockGeminiClient, abortSignal);
expect(result).toBeNull();
});
});
//# sourceMappingURL=nextSpeakerChecker.test.js.map