Feat/add gemini (#349)

This commit is contained in:
Drew Payment
2024-07-04 04:03:54 -04:00
committed by GitHub
parent a4b4e65011
commit c58e0c62a4
20 changed files with 9265 additions and 7174 deletions
-1
View File
@@ -5,7 +5,6 @@ import { prepareEnvironment } from './utils';
it('cli flow when there are no changes', async () => {
const { gitDir, cleanup } = await prepareEnvironment();
const { findByText } = await render(`OCO_AI_PROVIDER='test' node`, [resolve('./out/cli.cjs')], { cwd: gitDir });
expect(await findByText('No changes detected')).toBeInTheConsole();
-1
View File
@@ -10,7 +10,6 @@ it('cli flow to generate commit message for 1 new file (staged)', async () => {
await render('git' ,['add index.ts'], { cwd: gitDir });
const { queryByText, findByText, userEvent } = await render(`OCO_AI_PROVIDER='test' node`, [resolve('./out/cli.cjs')], { cwd: gitDir });
expect(await queryByText('No files are staged')).not.toBeInTheConsole();
expect(await queryByText('Do you want to stage all files and generate commit message?')).not.toBeInTheConsole();
+3
View File
@@ -1,5 +1,8 @@
import 'cli-testing-library/extend-expect'
import { configure } from 'cli-testing-library'
import { jest } from '@jest/globals';
global.jest = jest;
/**
* Adjusted the wait time for waitFor/findByText to 2000ms, because the default 1000ms makes the test results flaky
+2 -2
View File
@@ -55,7 +55,7 @@ OCO_ONE_LINE_COMMIT="true"
expect(config!['OCO_LANGUAGE']).toEqual('de');
expect(config!['OCO_MESSAGE_TEMPLATE_PLACEHOLDER']).toEqual('$m');
expect(config!['OCO_PROMPT_MODULE']).toEqual('@commitlint');
expect(config!['OCO_AI_PROVIDER']).toEqual('ollama');
expect(() => ['ollama', 'gemini'].includes(config!['OCO_AI_PROVIDER'])).toBeTruthy();
expect(config!['OCO_GITPUSH']).toEqual(false);
expect(config!['OCO_ONE_LINE_COMMIT']).toEqual(true);
@@ -96,7 +96,7 @@ OCO_ONE_LINE_COMMIT="true"
expect(config!['OCO_LANGUAGE']).toEqual('de');
expect(config!['OCO_MESSAGE_TEMPLATE_PLACEHOLDER']).toEqual('$m');
expect(config!['OCO_PROMPT_MODULE']).toEqual('@commitlint');
expect(config!['OCO_AI_PROVIDER']).toEqual('ollama');
expect(() => ['ollama', 'gemini'].includes(config!['OCO_AI_PROVIDER'])).toBeTruthy();
expect(config!['OCO_GITPUSH']).toEqual(false);
expect(config!['OCO_ONE_LINE_COMMIT']).toEqual(true);
+105
View File
@@ -0,0 +1,105 @@
import { Gemini } from '../../src/engine/gemini';
import { ChatCompletionRequestMessage } from 'openai';
import { GenerativeModel, GoogleGenerativeAI } from '@google/generative-ai';
import { ConfigType, getConfig } from '../../src/commands/config';
describe('Gemini', () => {
let gemini: Gemini;
let mockConfig: ConfigType;
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 mockGemini = () => {
gemini = new Gemini();
}
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(),
}));
if (mockWarmup) mockWarmup.mockRestore();
mockExit = jest.spyOn(process, 'exit').mockImplementation();
mockConfig = getConfig() as ConfigType;
mockConfig.OCO_AI_PROVIDER = '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, });
});
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', () => {
process.env.OCO_GEMINI_API_KEY = undefined;
process.env.OCO_AI_PROVIDER = 'gemini';
mockGemini();
expect(mockExit).toHaveBeenCalledWith(1);
});
it('should generate commit message', async () => {
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 result = await gemini.generateCommitMessage(messages);
expect(result).toEqual('generated content');
expect(mockWarmup).toHaveBeenCalled();
});
});