Roo-Code/src/utils/__tests__/enhance-prompt.test.ts

import { enhancePrompt } from '../enhance-prompt'
import { ApiConfiguration } from '../../shared/api'
import { buildApiHandler, SingleCompletionHandler } from '../../api'
import { defaultPrompts } from '../../shared/modes'

// Mock the API handler
jest.mock('../../api', () => ({
  buildApiHandler: jest.fn()
}))

describe('enhancePrompt', () => {
  const mockApiConfig: ApiConfiguration = {
    apiProvider: 'openai',
    openAiApiKey: 'test-key',
    openAiBaseUrl: 'https://api.openai.com/v1'
  }

  beforeEach(() => {
    jest.clearAllMocks()
    
    // Mock the API handler with a completePrompt method
    ;(buildApiHandler as jest.Mock).mockReturnValue({
      completePrompt: jest.fn().mockResolvedValue('Enhanced prompt'),
      createMessage: jest.fn(),
      getModel: jest.fn().mockReturnValue({ 
        id: 'test-model',
        info: {
          maxTokens: 4096,
          contextWindow: 8192,
          supportsPromptCache: false
        }
      })
    } as unknown as SingleCompletionHandler)
  })

  it('enhances prompt using default enhancement prompt when no custom prompt provided', async () => {
    const result = await enhancePrompt(mockApiConfig, 'Test prompt')
    
    expect(result).toBe('Enhanced prompt')
    const handler = buildApiHandler(mockApiConfig)
    expect((handler as any).completePrompt).toHaveBeenCalledWith(
      `${defaultPrompts.enhance}\n\nTest prompt`
    )
  })

  it('enhances prompt using custom enhancement prompt when provided', async () => {
    const customEnhancePrompt = 'You are a custom prompt enhancer'
    
    const result = await enhancePrompt(mockApiConfig, 'Test prompt', customEnhancePrompt)
    
    expect(result).toBe('Enhanced prompt')
    const handler = buildApiHandler(mockApiConfig)
    expect((handler as any).completePrompt).toHaveBeenCalledWith(
      `${customEnhancePrompt}\n\nTest prompt`
    )
  })

  it('throws error for empty prompt input', async () => {
    await expect(enhancePrompt(mockApiConfig, '')).rejects.toThrow('No prompt text provided')
  })

  it('throws error for missing API configuration', async () => {
    await expect(enhancePrompt({} as ApiConfiguration, 'Test prompt')).rejects.toThrow('No valid API configuration provided')
  })

  it('throws error for API provider that does not support prompt enhancement', async () => {
    (buildApiHandler as jest.Mock).mockReturnValue({
      // No completePrompt method
      createMessage: jest.fn(),
      getModel: jest.fn().mockReturnValue({ 
        id: 'test-model',
        info: {
          maxTokens: 4096,
          contextWindow: 8192,
          supportsPromptCache: false
        }
      })
    })

    await expect(enhancePrompt(mockApiConfig, 'Test prompt')).rejects.toThrow('The selected API provider does not support prompt enhancement')
  })

  it('uses appropriate model based on provider', async () => {
    const openRouterConfig: ApiConfiguration = {
      apiProvider: 'openrouter',
      openRouterApiKey: 'test-key',
      openRouterModelId: 'test-model'
    }

    // Mock successful enhancement
    ;(buildApiHandler as jest.Mock).mockReturnValue({
      completePrompt: jest.fn().mockResolvedValue('Enhanced prompt'),
      createMessage: jest.fn(),
      getModel: jest.fn().mockReturnValue({ 
        id: 'test-model',
        info: {
          maxTokens: 4096,
          contextWindow: 8192,
          supportsPromptCache: false
        }
      })
    } as unknown as SingleCompletionHandler)

    const result = await enhancePrompt(openRouterConfig, 'Test prompt')
    
    expect(buildApiHandler).toHaveBeenCalledWith(openRouterConfig)
    expect(result).toBe('Enhanced prompt')
  })

  it('propagates API errors', async () => {
    (buildApiHandler as jest.Mock).mockReturnValue({
      completePrompt: jest.fn().mockRejectedValue(new Error('API Error')),
      createMessage: jest.fn(),
      getModel: jest.fn().mockReturnValue({ 
        id: 'test-model',
        info: {
          maxTokens: 4096,
          contextWindow: 8192,
          supportsPromptCache: false
        }
      })
    } as unknown as SingleCompletionHandler)

    await expect(enhancePrompt(mockApiConfig, 'Test prompt')).rejects.toThrow('API Error')
  })
})