Adding some simple tests. (#54)

2 files changed, 174 insertions, 0 deletions

diff --git a/packages/cli/src/core/gemini-client.test.ts b/packages/cli/src/core/gemini-client.test.ts
new file mode 100644
index 00000000..4651529e
--- /dev/null
+++ b/packages/cli/src/core/gemini-client.test.ts
@@ -0,0 +1,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
+});
diff --git a/packages/cli/src/index.test.ts b/packages/cli/src/index.test.ts
new file mode 100644
index 00000000..bbe9fa8b
--- /dev/null
+++ b/packages/cli/src/index.test.ts
@@ -0,0 +1,9 @@
+import { describe, it, expect } from 'vitest';
+import { toolRegistry } from './tools/tool-registry.js';
+
+describe('cli tests', () => {
+  it('should have a tool registry', () => {
+    expect(toolRegistry).toBeDefined();
+    expect(typeof toolRegistry.registerTool).toBe('function');
+  });
+});