lamplighter-mcp/__tests__/services/aiService.test.ts

An intelligent context engine for AI-assisted software development

import 'openai/shims/node'; // Add node fetch shim
import { AIService, IAIService, AIConfig } from '../../src/services/aiService';
import OpenAI from 'openai';

// --- Mocks ---

// Mock OpenAI
jest.mock('openai');
const mockedOpenAI = jest.mocked(OpenAI);
const mockCreate = jest.fn();
mockedOpenAI.mockImplementation(() => ({
    chat: {
        completions: {
            create: mockCreate,
        },
    },
} as unknown as OpenAI));

// Mock fetch for Google Gemini
global.fetch = jest.fn();
const mockedFetch = jest.mocked(global.fetch);

// --- Environment Variable Management ---
const originalEnv = process.env;

beforeEach(() => {
  // Reset mocks and environment variables before each test
  jest.clearAllMocks();
  process.env = { ...originalEnv }; // Start with original env
  // Set default required keys for tests, can be overridden
  process.env.OPENAI_API_KEY = 'test-openai-key';
  process.env.GOOGLE_API_KEY = 'test-google-key';
});

afterAll(() => {
  // Restore original environment variables
  process.env = originalEnv;
});

// --- Test Suites ---

describe('AIService Factory (getInstance)', () => {
  beforeEach(() => {
    // Reset the singleton and mocks before each test
    (AIService as any).instance = undefined;
    mockedOpenAI.mockClear();
  });

  it('should return OpenAIService by default', () => {
    delete process.env.AI_PROVIDER;
    const instance = AIService.getInstance();
    // Check if the instance is an OpenAIService (difficult directly, check constructor side effect)
    expect(mockedOpenAI).toHaveBeenCalled();
  });

  it('should return OpenAIService when AI_PROVIDER is openai', () => {
    process.env.AI_PROVIDER = 'openai';
    const instance = AIService.getInstance();
    expect(mockedOpenAI).toHaveBeenCalled();
  });

  it('should return GoogleGeminiService when AI_PROVIDER is google', () => {
    process.env.AI_PROVIDER = 'google';
    const instance = AIService.getInstance();
    // Check constructor side-effect (no direct OpenAI mock call)
    expect(mockedOpenAI).not.toHaveBeenCalled();
    // Check if it implements IAIService (basic check)
    expect(instance.generateText).toBeDefined();
  });

  it('should return the same instance on subsequent calls (singleton)', () => {
    process.env.AI_PROVIDER = 'google';
    const instance1 = AIService.getInstance();
    const instance2 = AIService.getInstance();
    expect(instance1).toBe(instance2);
  });

  it('should re-initialize if instance is manually reset (for testing purposes)', () => {
    // This test demonstrates how to reset the singleton if needed between test suites
    mockedOpenAI.mockClear(); // Clear mock counts first
    process.env.AI_PROVIDER = 'openai';
    const instance1 = AIService.getInstance();
    expect(mockedOpenAI).toHaveBeenCalledTimes(1);

    // Manually reset internal state
    (AIService as any).instance = undefined;

    process.env.AI_PROVIDER = 'google';
    const instance2 = AIService.getInstance();
    expect(mockedOpenAI).toHaveBeenCalledTimes(1); // OpenAI should not be called again
    expect(instance1).not.toBe(instance2);
  });
});

describe('OpenAIService', () => {
  let openAIService: IAIService;

  beforeEach(() => {
    // Ensure we get an OpenAIService instance
    (AIService as any).instance = undefined; // Reset singleton
    process.env.AI_PROVIDER = 'openai';
    openAIService = AIService.getInstance();
    // Reset OpenAI mock function calls specifically for this suite
    mockCreate.mockClear();
  });

  it('should throw if OPENAI_API_KEY is missing', () => {
    delete process.env.OPENAI_API_KEY;
    (AIService as any).instance = undefined; // Reset to force re-creation
    expect(() => AIService.getInstance()).toThrow('OpenAI API key is required');
  });

  it('should call OpenAI create with correct parameters', async () => {
    const prompt = 'Test prompt';
    const config: AIConfig = { temperature: 0.5, maxTokens: 100 };
    const mockResponse = { choices: [{ message: { content: 'OpenAI response' } }] };
    mockCreate.mockResolvedValue(mockResponse);

    await openAIService.generateText(prompt, config);

    expect(mockCreate).toHaveBeenCalledTimes(1);
    expect(mockCreate).toHaveBeenCalledWith({
      model: process.env.AI_MODEL || 'gpt-4', // Check default or env model
      messages: [{ role: 'user', content: prompt }],
      temperature: config.temperature,
      max_tokens: config.maxTokens,
      top_p: 1, // Default top_p
    });
  });

   it('should use default config values if not provided', async () => {
    const prompt = 'Test prompt defaults';
    const mockResponse = { choices: [{ message: { content: 'Default response' } }] };
    mockCreate.mockResolvedValue(mockResponse);

    await openAIService.generateText(prompt); // No config provided

    expect(mockCreate).toHaveBeenCalledTimes(1);
    expect(mockCreate).toHaveBeenCalledWith({
      model: process.env.AI_MODEL || 'gpt-4',
      messages: [{ role: 'user', content: prompt }],
      temperature: 0.7, // Default temperature
      max_tokens: undefined, // Default max_tokens
      top_p: 1, // Default top_p
    });
  });

  it('should return the content from the response', async () => {
    const expectedContent = 'Successful OpenAI response';
    const mockResponse = { choices: [{ message: { content: expectedContent } }] };
    mockCreate.mockResolvedValue(mockResponse);

    const result = await openAIService.generateText('Get response');
    expect(result).toBe(expectedContent);
  });

  it('should handle OpenAI API errors', async () => {
    const apiError = new Error('Invalid API Key');
    mockCreate.mockRejectedValue(apiError);

    await expect(openAIService.generateText('Trigger error')).rejects.toThrow(
      `Failed to generate text with OpenAI: ${apiError.message}`
    );
  });
});

describe('GoogleGeminiService', () => {
  let googleService: IAIService;

  beforeEach(() => {
    // Ensure we get a GoogleGeminiService instance
    (AIService as any).instance = undefined; // Reset singleton
    process.env.AI_PROVIDER = 'google';
    googleService = AIService.getInstance();
    // Reset fetch mock specifically for this suite
    mockedFetch.mockClear();
  });

  it('should throw if GOOGLE_API_KEY is missing', () => {
    delete process.env.GOOGLE_API_KEY;
    (AIService as any).instance = undefined; // Reset to force re-creation
    expect(() => AIService.getInstance()).toThrow('Google API key is required');
  });

  it('should call fetch with correct parameters', async () => {
    const prompt = 'Google test prompt';
    const config: AIConfig = { temperature: 0.6, maxTokens: 150, topP: 0.9 };
    const expectedContent = 'Google response';
    const mockResponse = { candidates: [{ content: { parts: [{ text: expectedContent }] } }] };
    mockedFetch.mockResolvedValue({
      ok: true,
      json: async () => mockResponse,
      text: async () => JSON.stringify(mockResponse),
    } as Response);

    await googleService.generateText(prompt, config);

    expect(mockedFetch).toHaveBeenCalledTimes(1);
    const fetchCall = mockedFetch.mock.calls[0];
    const url = fetchCall[0];
    const options = fetchCall[1];

    const expectedModel = process.env.AI_MODEL || 'gemini-pro';
    expect(url).toContain(`https://generativelanguage.googleapis.com/v1/models/${expectedModel}:generateContent`);
    expect(url).toContain(`key=${process.env.GOOGLE_API_KEY}`);

    expect(options?.method).toBe('POST');
    expect(options?.headers).toEqual({ 'Content-Type': 'application/json' });
    expect(JSON.parse(options?.body as string)).toEqual({
      contents: [{ parts: [{ text: prompt }] }],
      generationConfig: {
        temperature: config.temperature,
        topP: config.topP,
        maxOutputTokens: config.maxTokens,
      },
    });
  });

  it('should use default config values if not provided', async () => {
    const prompt = 'Google defaults';
    const mockResponse = { candidates: [{ content: { parts: [{ text: 'Default google' }] } }] };
    mockedFetch.mockResolvedValue({
      ok: true,
      json: async () => mockResponse,
       text: async () => JSON.stringify(mockResponse),
    } as Response);

    await googleService.generateText(prompt); // No config

    expect(mockedFetch).toHaveBeenCalledTimes(1);
    const options = mockedFetch.mock.calls[0][1];
    expect(JSON.parse(options?.body as string).generationConfig).toEqual({
      temperature: 0.7, // Default temperature
      topP: 1, // Default topP
      maxOutputTokens: 1024, // Default maxOutputTokens
    });
  });

  it('should return the content from the response', async () => {
    const expectedContent = 'Successful Google response';
    const mockResponse = { candidates: [{ content: { parts: [{ text: expectedContent }] } }] };
    mockedFetch.mockResolvedValue({
      ok: true,
      json: async () => mockResponse,
       text: async () => JSON.stringify(mockResponse),
    } as Response);

    const result = await googleService.generateText('Get response');
    expect(result).toBe(expectedContent);
  });

  it('should handle Google API errors (non-OK response)', async () => {
    const errorText = 'Invalid API key or resource exhausted.';
    mockedFetch.mockResolvedValue({
      ok: false,
      status: 400,
      json: async () => ({ error: { message: errorText } }),
      text: async () => errorText,
    } as Response);

    await expect(googleService.generateText('Trigger error')).rejects.toThrow(
      `Failed to generate text with Google Gemini: Google API error: 400 ${errorText}`
    );
  });

  it('should handle network errors during fetch', async () => {
    const networkError = new Error('Fetch failed');
    mockedFetch.mockRejectedValue(networkError);

    await expect(googleService.generateText('Trigger network error')).rejects.toThrow(
      `Failed to generate text with Google Gemini: ${networkError.message}`
    );
  });
});