konnosDashboard
generate-v4.test.ts86 lines · main
1import { safeSql } from '@supabase/pg-meta'
2import { UIMessage } from 'ai'
3import { expect, test, vi } from 'vitest'
4
5import generateV4 from '../../pages/api/ai/sql/generate-v4'
6import { sanitizeMessagePart } from '@/lib/ai/tools/tool-sanitizer'
7
8vi.mock('@/lib/ai/tools/tool-sanitizer', () => ({
9 sanitizeMessagePart: vi.fn((part) => part),
10}))
11
12test('generateV4 calls the tool sanitizer', async () => {
13 const mockReq = {
14 method: 'POST',
15 headers: {
16 authorization: 'Bearer test-token',
17 },
18 body: {
19 messages: [
20 {
21 id: 'test-msg-id',
22 role: 'assistant',
23 parts: [
24 {
25 type: 'tool-execute_sql',
26 state: 'output-available',
27 toolCallId: 'test-tool-call-id',
28 input: { sql: safeSql`SELECT * FROM users` },
29 output: [{ id: 1, name: 'test-output' }],
30 },
31 ],
32 },
33 ] satisfies UIMessage[],
34 projectRef: 'test-project',
35 connectionString: 'test-connection',
36 orgSlug: 'test-org',
37 },
38 on: vi.fn(),
39 }
40
41 const mockRes = {
42 status: vi.fn(() => mockRes),
43 json: vi.fn(() => mockRes),
44 setHeader: vi.fn(() => mockRes),
45 }
46
47 vi.mock('@/lib/ai/ai-details', () => ({
48 getOrgAIDetails: vi.fn().mockResolvedValue({
49 aiOptInLevel: 'schema_and_log_and_data',
50 hasAccessToAdvanceModel: true,
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 systemProviderOptions: {},
62 }),
63 }))
64
65 vi.mock('@/data/sql/execute-sql-query', () => ({
66 executeSql: vi.fn().mockResolvedValue({ result: [] }),
67 }))
68
69 vi.mock('@/lib/ai/tools', () => ({
70 getTools: vi.fn().mockResolvedValue({}),
71 }))
72
73 vi.mock('ai', async () => {
74 const actual = await vi.importActual('ai')
75 return {
76 ...actual,
77 streamText: vi.fn().mockReturnValue({
78 pipeUIMessageStreamToResponse: vi.fn(),
79 }),
80 }
81 })
82
83 await generateV4(mockReq as any, mockRes as any)
84
85 expect(sanitizeMessagePart).toHaveBeenCalled()
86})