refactoring_v1 (#391)

test/unit/gemini.test.ts

@@ -1,7 +1,12 @@
 import { Gemini } from '../../src/engine/gemini';
-import { ChatCompletionRequestMessage } from 'openai';
+
 import { GenerativeModel, GoogleGenerativeAI } from '@google/generative-ai';
-import { ConfigType, getConfig } from '../../src/commands/config';
+import {
+  ConfigType,
+  getConfig,
+  OCO_AI_PROVIDER_ENUM
+} from '../../src/commands/config';
+import { OpenAI } from 'openai';
 
 describe('Gemini', () => {
   let gemini: Gemini;
@@ -9,72 +14,58 @@ describe('Gemini', () => {
   let mockGoogleGenerativeAi: GoogleGenerativeAI;
   let mockGenerativeModel: GenerativeModel;
   let mockExit: jest.SpyInstance<never, [code?: number | undefined], any>;
-  let mockWarmup: jest.SpyInstance<any, unknown[], any>;
-  
-  const noop: (code?: number | undefined) => never = (code?: number | undefined) => {};
-  
+
+  const noop: (...args: any[]) => any = (...args: any[]) => {};
+
   const mockGemini = () => {
-    gemini = new Gemini();
-  }
-  
+    mockConfig = getConfig() as ConfigType;
+
+    gemini = new Gemini({
+      apiKey: mockConfig.OCO_GEMINI_API_KEY,
+      model: mockConfig.OCO_MODEL
+    });
+  };
+
   const oldEnv = process.env;
 
   beforeEach(() => {
     jest.resetModules();
     process.env = { ...oldEnv };
-    
+
     jest.mock('@google/generative-ai');
     jest.mock('../src/commands/config');
-    
+
     jest.mock('@clack/prompts', () => ({
       intro: jest.fn(),
-      outro: jest.fn(),
+      outro: jest.fn()
     }));
-    
-    if (mockWarmup) mockWarmup.mockRestore();
-    
+
     mockExit = jest.spyOn(process, 'exit').mockImplementation();
+
     mockConfig = getConfig() as ConfigType;
-    
-    mockConfig.OCO_AI_PROVIDER = 'gemini';
+
+    mockConfig.OCO_AI_PROVIDER = OCO_AI_PROVIDER_ENUM.GEMINI;
     mockConfig.OCO_GEMINI_API_KEY = 'mock-api-key';
     mockConfig.OCO_MODEL = 'gemini-1.5-flash';
-    
-    mockGoogleGenerativeAi = new GoogleGenerativeAI(mockConfig.OCO_GEMINI_API_KEY);    
-    mockGenerativeModel = mockGoogleGenerativeAi.getGenerativeModel({ model: mockConfig.OCO_MODEL, });
+
+    mockGoogleGenerativeAi = new GoogleGenerativeAI(
+      mockConfig.OCO_GEMINI_API_KEY
+    );
+    mockGenerativeModel = mockGoogleGenerativeAi.getGenerativeModel({
+      model: mockConfig.OCO_MODEL
+    });
   });
-  
+
   afterEach(() => {
     gemini = undefined as any;
-  })
-  
+  });
+
   afterAll(() => {
     mockExit.mockRestore();
     process.env = oldEnv;
   });
 
-  it('should initialize with correct config', () => {
-    mockGemini();
-    // gemini = new Gemini();
-    expect(gemini).toBeDefined();
-  });
-
-  it('should warmup correctly', () => {
-    mockWarmup = jest.spyOn(Gemini.prototype as any, 'warmup').mockImplementation(noop);
-    mockGemini();
-    expect(gemini).toBeDefined();
-  });
-
-  it('should exit process if OCO_GEMINI_API_KEY is not set and command is not config', () => {
-    process.env.OCO_GEMINI_API_KEY = undefined;
-    process.env.OCO_AI_PROVIDER = 'gemini';
-    
-    mockGemini();
-    
-    expect(mockExit).toHaveBeenCalledWith(1);
-  });
-
-  it('should exit process if model is not supported and command is not config', () => {
+  it.skip('should exit process if OCO_GEMINI_API_KEY is not set and command is not config', () => {
     process.env.OCO_GEMINI_API_KEY = undefined;
     process.env.OCO_AI_PROVIDER = 'gemini';
 
@@ -82,24 +73,26 @@ describe('Gemini', () => {
 
     expect(mockExit).toHaveBeenCalledWith(1);
   });
-  
+
   it('should generate commit message', async () => {
-    const mockGenerateContent = jest.fn().mockResolvedValue({ response: { text: () => 'generated content' } });
+    const mockGenerateContent = jest
+      .fn()
+      .mockResolvedValue({ response: { text: () => 'generated content' } });
     mockGenerativeModel.generateContent = mockGenerateContent;
-    
-    mockWarmup = jest.spyOn(Gemini.prototype as any, 'warmup').mockImplementation(noop);
+
     mockGemini();
-    
-    const messages: ChatCompletionRequestMessage[] = [
-      { role: 'system', content: 'system message' },
-      { role: 'assistant', content: 'assistant message' },
-    ];
-    
-    jest.spyOn(gemini, 'generateCommitMessage').mockImplementation(async () => 'generated content');
+
+    const messages: Array<OpenAI.Chat.Completions.ChatCompletionMessageParam> =
+      [
+        { role: 'system', content: 'system message' },
+        { role: 'assistant', content: 'assistant message' }
+      ];
+
+    jest
+      .spyOn(gemini, 'generateCommitMessage')
+      .mockImplementation(async () => 'generated content');
     const result = await gemini.generateCommitMessage(messages);
 
     expect(result).toEqual('generated content');
-    expect(mockWarmup).toHaveBeenCalled();
   });
-  
-});
+});