Add a screen for custom prompts

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

@@ -1,80 +1,126 @@
 import { enhancePrompt } from '../enhance-prompt'
-import { buildApiHandler } from '../../api'
 import { ApiConfiguration } from '../../shared/api'
-import { OpenRouterHandler } from '../../api/providers/openrouter'
+import { buildApiHandler, SingleCompletionHandler } from '../../api'
+import { defaultPrompts } from '../../shared/modes'
 
-// Mock the buildApiHandler function
+// Mock the API handler
 jest.mock('../../api', () => ({
-    buildApiHandler: jest.fn()
+  buildApiHandler: jest.fn()
 }))
 
 describe('enhancePrompt', () => {
-    const mockApiConfig: ApiConfiguration = {
-        apiProvider: 'openrouter',
-        apiKey: 'test-key',
-        openRouterApiKey: 'test-key',
-        openRouterModelId: 'test-model'
-    }
+  const mockApiConfig: ApiConfiguration = {
+    apiProvider: 'openai',
+    openAiApiKey: 'test-key',
+    openAiBaseUrl: 'https://api.openai.com/v1'
+  }
 
-    // Create a mock handler that looks like OpenRouterHandler
-    const mockHandler = {
-        completePrompt: jest.fn(),
-        createMessage: jest.fn(),
-        getModel: jest.fn()
-    }
-
-    // Make instanceof check work
-    Object.setPrototypeOf(mockHandler, OpenRouterHandler.prototype)
-
-    beforeEach(() => {
-        jest.clearAllMocks()
-        ;(buildApiHandler as jest.Mock).mockReturnValue(mockHandler)
-    })
-
-    it('should throw error for non-OpenRouter providers', async () => {
-        const nonOpenRouterConfig: ApiConfiguration = {
-            apiProvider: 'anthropic',
-            apiKey: 'test-key',
-            apiModelId: 'claude-3'
+  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
         }
-        await expect(enhancePrompt(nonOpenRouterConfig, 'test')).rejects.toThrow('Prompt enhancement is only available with OpenRouter')
+      })
+    } 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
+        }
+      })
     })
 
-    it('should enhance a valid prompt', async () => {
-        const inputPrompt = 'Write a function to sort an array'
-        const enhancedPrompt = 'Write a TypeScript function that implements an efficient sorting algorithm for a generic array, including error handling and type safety'
-        
-        mockHandler.completePrompt.mockResolvedValue(enhancedPrompt)
+    await expect(enhancePrompt(mockApiConfig, 'Test prompt')).rejects.toThrow('The selected API provider does not support prompt enhancement')
+  })
 
-        const result = await enhancePrompt(mockApiConfig, inputPrompt)
+  it('uses appropriate model based on provider', async () => {
+    const openRouterConfig: ApiConfiguration = {
+      apiProvider: 'openrouter',
+      openRouterApiKey: 'test-key',
+      openRouterModelId: 'test-model'
+    }
 
-        expect(result).toBe(enhancedPrompt)
-        expect(buildApiHandler).toHaveBeenCalledWith(mockApiConfig)
-        expect(mockHandler.completePrompt).toHaveBeenCalledWith(
-            expect.stringContaining(inputPrompt)
-        )
-    })
+    // 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)
 
-    it('should throw error when no prompt text is provided', async () => {
-        await expect(enhancePrompt(mockApiConfig, '')).rejects.toThrow('No prompt text provided')
-        expect(mockHandler.completePrompt).not.toHaveBeenCalled()
-    })
+    const result = await enhancePrompt(openRouterConfig, 'Test prompt')
+    
+    expect(buildApiHandler).toHaveBeenCalledWith(openRouterConfig)
+    expect(result).toBe('Enhanced prompt')
+  })
 
-    it('should pass through API errors', async () => {
-        const inputPrompt = 'Test prompt'
-        mockHandler.completePrompt.mockRejectedValue('API error')
+  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, inputPrompt)).rejects.toBe('API error')
-    })
-
-    it('should pass the correct prompt format to the API', async () => {
-        const inputPrompt = 'Test prompt'
-        mockHandler.completePrompt.mockResolvedValue('Enhanced test prompt')
-        
-        await enhancePrompt(mockApiConfig, inputPrompt)
-
-        expect(mockHandler.completePrompt).toHaveBeenCalledWith(
-            'Generate an enhanced version of this prompt (reply with only the enhanced prompt - no conversation, explanations, lead-in, bullet points, placeholders, or surrounding quotes):\n\nTest prompt'
-        )
-    })
+    await expect(enhancePrompt(mockApiConfig, 'Test prompt')).rejects.toThrow('API Error')
+  })
 })