konnosDashboard
rate.test.ts80 lines · main
1import { expect, test, vi } from 'vitest'
2
3// End of third-party imports
4
5import rate from '../../pages/api/ai/feedback/rate'
6import { sanitizeMessagePart } from '../ai/tools/tool-sanitizer'
7
8vi.mock('../ai/tools/tool-sanitizer', () => ({
9 sanitizeMessagePart: vi.fn((part) => part),
10}))
11
12test('rate calls the tool sanitizer', async () => {
13 const mockReq = {
14 method: 'POST',
15 headers: {
16 authorization: 'Bearer test-token',
17 },
18 body: {
19 rating: 'negative',
20 messages: [
21 {
22 role: 'assistant',
23 parts: [
24 {
25 type: 'tool-execute_sql',
26 state: 'output-available',
27 output: 'test output',
28 },
29 ],
30 },
31 ],
32 messageId: 'test-message-id',
33 projectRef: 'test-project',
34 orgSlug: 'test-org',
35 reason: 'The response was not helpful',
36 },
37 on: vi.fn(),
38 }
39
40 const mockRes = {
41 status: vi.fn(() => mockRes),
42 json: vi.fn(() => mockRes),
43 setHeader: vi.fn(() => mockRes),
44 }
45
46 vi.mock('@/lib/ai/ai-details', () => ({
47 getOrgAIDetails: vi.fn().mockResolvedValue({
48 aiOptInLevel: 'schema_and_log_and_data',
49 hasAccessToAdvanceModel: true,
50 isDpaSigned: false,
51 }),
52 getProjectAIDetails: vi.fn().mockResolvedValue({
53 region: 'us-east-1',
54 isSensitive: false,
55 }),
56 }))
57
58 vi.mock('@/lib/ai/model', () => ({
59 getModel: vi.fn().mockResolvedValue({
60 modelParams: { model: {} },
61 }),
62 }))
63
64 vi.mock('ai', () => ({
65 generateText: vi.fn().mockResolvedValue({
66 output: {
67 category: 'sql_generation',
68 },
69 }),
70 Output: { object: vi.fn() },
71 }))
72
73 vi.mock('@/components/ui/AIAssistantPanel/Message.utils', () => ({
74 rateMessageResponseSchema: {},
75 }))
76
77 await rate(mockReq as any, mockRes as any)
78
79 expect(sanitizeMessagePart).toHaveBeenCalled()
80})