1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
|
import { describe, it, expect, vi, beforeEach, Mock } from 'vitest';
import { GoogleGenAI, Type, Content } from '@google/genai';
import { GeminiClient } from './gemini-client.js';
import { Config } from '../config/config.js';
// Mock the entire @google/genai module
vi.mock('@google/genai');
// Mock the Config class and its methods
vi.mock('../config/config.js', () => {
// The mock constructor should accept the arguments but not explicitly return an object.
// vi.fn() will create a mock instance that inherits from the prototype.
const MockConfig = vi.fn();
// Methods are mocked on the prototype, so instances will inherit them.
MockConfig.prototype.getApiKey = vi.fn(() => 'mock-api-key');
MockConfig.prototype.getModel = vi.fn(() => 'mock-model');
MockConfig.prototype.getTargetDir = vi.fn(() => 'mock-target-dir');
return { Config: MockConfig };
});
// Define a type for the mocked GoogleGenAI instance structure
type MockGoogleGenAIType = {
models: {
generateContent: Mock;
};
chats: {
create: Mock;
};
};
describe('GeminiClient', () => {
// Use the specific types defined above
let mockGenerateContent: MockGoogleGenAIType['models']['generateContent'];
let mockGoogleGenAIInstance: MockGoogleGenAIType;
let config: Config;
let client: GeminiClient;
beforeEach(() => {
vi.clearAllMocks();
// Mock the generateContent method specifically
mockGenerateContent = vi.fn();
// Mock the chainable structure ai.models.generateContent
mockGoogleGenAIInstance = {
models: {
generateContent: mockGenerateContent,
},
chats: {
create: vi.fn(), // Mock create as well
},
};
// Configure the mocked GoogleGenAI constructor to return our mock instance
(GoogleGenAI as Mock).mockImplementation(() => mockGoogleGenAIInstance);
config = new Config('mock-api-key-arg', 'mock-model-arg', 'mock-dir-arg');
client = new GeminiClient(config);
});
describe('generateJson', () => {
it('should call ai.models.generateContent with correct parameters', async () => {
const mockContents: Content[] = [
{ role: 'user', parts: [{ text: 'test prompt' }] },
];
const mockSchema = {
type: Type.OBJECT,
properties: { key: { type: Type.STRING } },
};
const mockApiResponse = { text: JSON.stringify({ key: 'value' }) };
mockGenerateContent.mockResolvedValue(mockApiResponse);
await client.generateJson(mockContents, mockSchema);
expect(mockGenerateContent).toHaveBeenCalledTimes(1);
// Use expect.objectContaining for the config assertion
const expectedConfigMatcher = expect.objectContaining({
temperature: 0,
topP: 1,
systemInstruction: expect.any(String),
responseSchema: mockSchema,
responseMimeType: 'application/json',
});
expect(mockGenerateContent).toHaveBeenCalledWith({
model: 'mock-model',
config: expectedConfigMatcher,
contents: mockContents,
});
});
it('should return the parsed JSON response', async () => {
const mockContents: Content[] = [
{ role: 'user', parts: [{ text: 'test prompt' }] },
];
const mockSchema = {
type: Type.OBJECT,
properties: { key: { type: Type.STRING } },
};
const expectedJson = { key: 'value' };
const mockApiResponse = { text: JSON.stringify(expectedJson) };
mockGenerateContent.mockResolvedValue(mockApiResponse);
const result = await client.generateJson(mockContents, mockSchema);
expect(result).toEqual(expectedJson);
});
it('should throw an error if API returns empty response', async () => {
const mockContents: Content[] = [
{ role: 'user', parts: [{ text: 'test prompt' }] },
];
const mockSchema = {
type: Type.OBJECT,
properties: { key: { type: Type.STRING } },
};
const mockApiResponse = { text: '' }; // Empty response
mockGenerateContent.mockResolvedValue(mockApiResponse);
await expect(
client.generateJson(mockContents, mockSchema),
).rejects.toThrow(
'Failed to generate JSON content: API returned an empty response.',
);
});
it('should throw an error if API response is not valid JSON', async () => {
const mockContents: Content[] = [
{ role: 'user', parts: [{ text: 'test prompt' }] },
];
const mockSchema = {
type: Type.OBJECT,
properties: { key: { type: Type.STRING } },
};
const mockApiResponse = { text: 'invalid json' }; // Invalid JSON
mockGenerateContent.mockResolvedValue(mockApiResponse);
await expect(
client.generateJson(mockContents, mockSchema),
).rejects.toThrow('Failed to parse API response as JSON:');
});
it('should throw an error if generateContent rejects', async () => {
const mockContents: Content[] = [
{ role: 'user', parts: [{ text: 'test prompt' }] },
];
const mockSchema = {
type: Type.OBJECT,
properties: { key: { type: Type.STRING } },
};
const apiError = new Error('API call failed');
mockGenerateContent.mockRejectedValue(apiError);
await expect(
client.generateJson(mockContents, mockSchema),
).rejects.toThrow(`Failed to generate JSON content: ${apiError.message}`);
});
});
// TODO: Add tests for startChat and sendMessageStream later
});
|
