import { untrustedSql } from '@supabase/pg-meta' import type { QueryClient } from '@tanstack/react-query' import { renderHook, type RenderHookOptions } from '@testing-library/react' import { FeatureFlagContext } from 'common' import { http, HttpResponse } from 'msw' import { NuqsTestingAdapter } from 'nuqs/adapters/testing' import type { ReactNode } from 'react' import type { SqlSnippetSource } from '@/components/interfaces/SQLEditor/querySource' import type { ContentDiff, DiffController, EditorController, } from '@/components/interfaces/SQLEditor/SQLEditor.types' import { computeErrorHighlightLine, createSqlSnippetSkeletonV2, } from '@/components/interfaces/SQLEditor/SQLEditor.utils' import { SQLEditorProvider } from '@/components/interfaces/SQLEditor/SQLEditorContext' import { API_URL, OPT_IN_TAGS } from '@/lib/constants' import type { ProfileContextType } from '@/lib/profile' import { AiAssistantStateContext, createAiAssistantState } from '@/state/ai-assistant-state' import { createDatabaseSelectorState, DatabaseSelectorStateContext, } from '@/state/database-selector' import { createRoleImpersonationState, RoleImpersonationStateContext, } from '@/state/role-impersonation-state' import { sidebarManagerState } from '@/state/sidebar-manager-state' import { sqlEditorDiffRequestState } from '@/state/sql-editor/sql-editor-diff-request' import { sqlEditorSessionState } from '@/state/sql-editor/sql-editor-session-state' import { sqlEditorState } from '@/state/sql-editor/sql-editor-state' import { CustomWrapper } from '@/tests/lib/custom-render' import { addAPIMock, mswServer } from '@/tests/lib/msw' export type InMemoryEditor = { editor: EditorController diff: DiffController setValue: (value: string) => void setSelection: (value: string | undefined, startLineNumber?: number) => void getHighlightedLine: () => number | undefined getRevealedLine: () => number | undefined getDiffOriginal: () => string getDiffModified: () => string getDiffRevealedLine: () => number | undefined setDiffMounted: (mounted: boolean) => void } export function createInMemoryEditor(initialValue: string = ''): InMemoryEditor { let value = initialValue let selectionValue: string | undefined let selectionStartLine: number | undefined let highlightedLine: number | undefined let revealedLine: number | undefined let diffMounted = true let diffOriginal = '' let diffModified = '' let diffRevealedLine: number | undefined const editor: EditorController = { isReady: () => true, getValue: () => value, getSelectionStartLine: () => selectionStartLine, getSql: (snippetContent) => untrustedSql(selectionValue || value || snippetContent || ''), replaceAll: (text) => { value = text }, focus: () => {}, revealLineInCenter: (line) => { revealedLine = line }, highlightErrorLine: (error, hasSelection) => { const startLineNumber = hasSelection ? (selectionStartLine ?? 0) : 0 const line = computeErrorHighlightLine(error, startLineNumber) if (Number.isNaN(line)) return highlightedLine = line revealedLine = line }, clearHighlights: () => { highlightedLine = undefined }, } const diff: DiffController = { isMounted: () => diffMounted, getModifiedValue: () => diffModified, setDiff: (contentDiff: ContentDiff, revealLine: number) => { diffOriginal = contentDiff.original diffModified = contentDiff.modified diffRevealedLine = revealLine }, attach: () => { diffMounted = true }, } return { editor, diff, setValue: (v) => { value = v }, setSelection: (v, startLine = 1) => { selectionValue = v selectionStartLine = v === undefined ? undefined : startLine }, getHighlightedLine: () => highlightedLine, getRevealedLine: () => revealedLine, getDiffOriginal: () => diffOriginal, getDiffModified: () => diffModified, getDiffRevealedLine: () => diffRevealedLine, setDiffMounted: (mounted) => { diffMounted = mounted }, } } /** Call in `beforeEach` for full isolation between SQL editor hook tests. */ export function resetSqlEditorStores() { for (const key of Object.keys(sqlEditorState.snippets)) { delete sqlEditorState.snippets[key] } for (const key of Object.keys(sqlEditorState.folders)) { delete sqlEditorState.folders[key] } sqlEditorState.needsSaving.clear() sqlEditorState.pendingFolderSaves.clear() for (const key of Object.keys(sqlEditorSessionState.results)) { delete sqlEditorSessionState.results[key] } sqlEditorSessionState.limit = 100 sqlEditorDiffRequestState.pending = undefined sidebarManagerState.sidebars = {} sidebarManagerState.activeSidebar = undefined sidebarManagerState.pendingSidebarOpen = undefined sidebarManagerState.isMaximised = false } export function seedSnippet({ id, projectRef = 'default', name = 'Test query', sql = '', ownerId = 1, projectId = 1, source = 'database', }: { id: string projectRef?: string name?: string sql?: string ownerId?: number projectId?: number source?: SqlSnippetSource }) { const snippet = createSqlSnippetSkeletonV2({ name, sql, owner_id: ownerId, project_id: projectId, idOverride: id, source, }) sqlEditorState.addSnippet({ projectRef, snippet }) return snippet } export function setupSqlEditorMocks({ ref = 'default', connectionString = 'postgresql://postgres@localhost:5432/postgres', orgId = 1, orgSlug = 'test-org', optInTags = [OPT_IN_TAGS.AI_SQL], queryRows = [] as unknown[], }: { ref?: string connectionString?: string orgId?: number orgSlug?: string optInTags?: string[] queryRows?: unknown[] } = {}) { addAPIMock({ method: 'get', path: '/platform/projects/:ref', response: { id: 1, ref, organization_id: orgId, name: 'Test Project', status: 'ACTIVE_HEALTHY', cloud_provider: 'AWS', region: 'us-east-1', db_host: `db.${ref}.supabase.co`, restUrl: `https://${ref}.supabase.co/rest/v1/`, inserted_at: '2024-01-01T00:00:00Z', updated_at: '2024-01-01T00:00:00Z', subscription_id: 'sub_123', is_branch_enabled: false, is_physical_backups_enabled: false, high_availability: false, integration_source: null, connectionString, is_hibernating: false, }, }) addAPIMock({ method: 'get', path: '/platform/organizations', response: [ { id: orgId, slug: orgSlug, name: 'Test Org', billing_email: 'test@supabase.io', billing_partner: null, integration_source: null, is_owner: true, opt_in_tags: optInTags, organization_missing_address: false, organization_missing_tax_id: false, organization_requires_mfa: false, requires_indirect_tax_declaration: false, plan: { id: 'free', name: 'Free' }, restriction_data: null, restriction_status: null, stripe_customer_id: null, subscription_id: null, usage_billing_enabled: false, }, ], }) addAPIMock({ method: 'get', path: '/platform/organizations/:slug/billing/subscription', response: { addons: [], billing_cycle_anchor: 1700000000, billing_via_partner: false, current_period_end: 1700000000, current_period_start: 1700000000, next_invoice_at: 1700000000, payment_method_type: 'card', plan: { id: 'free', name: 'Free' }, project_addons: [], scheduled_plan_change: null, usage_billing_enabled: false, }, }) addAPIMock({ method: 'get', path: '/platform/projects/:ref/settings', response: { cloud_provider: 'AWS', db_dns_name: `db.${ref}.supabase.co`, db_host: `db.${ref}.supabase.co`, db_ip_addr_config: 'ipv4', db_name: 'postgres', db_port: 5432, db_user: 'postgres', inserted_at: '2024-01-01T00:00:00Z', name: 'Test Project', ref, region: 'us-east-1', ssl_enforced: false, status: 'ACTIVE_HEALTHY', }, }) addAPIMock({ method: 'get', path: '/platform/projects/:ref/databases', response: [ { identifier: ref, connectionString, cloud_provider: 'AWS', db_host: `db.${ref}.supabase.co`, db_name: 'postgres', db_port: 5432, db_user: 'postgres', inserted_at: '2024-01-01T00:00:00Z', region: 'us-east-1', restUrl: `https://${ref}.supabase.co/rest/v1/`, size: 'ci_micro', status: 'ACTIVE_HEALTHY', }, ], }) addAPIMock({ method: 'post', path: '/platform/pg-meta/:ref/query', response: () => HttpResponse.json(queryRows), }) // Internal Next API routes (not part of the platform OpenAPI spec), hit via // raw fetch — same MSW server, just registered directly rather than through // the typed `addAPIMock` helper. mswServer.use( http.post(`${API_URL}/ai/sql/title-v2`, async () => HttpResponse.json({ title: 'Generated title', description: '' }) ), http.post(`${API_URL}/ai/code/complete`, async () => HttpResponse.json('select 1;')) ) } type NuqsAdapterProps = Partial[0]> type RenderSqlEditorHookOptions = { initialProps?: TProps queryClient?: QueryClient nuqs?: NuqsAdapterProps profileContext?: ProfileContextType inMemoryEditor?: InMemoryEditor aiAssistantState?: ReturnType databaseSelectorState?: ReturnType roleImpersonationState?: ReturnType /** ConfigCat flags to expose via FeatureFlagContext (e.g. `{ otelLegacyLogs: true }`). */ flags?: Record } export function renderSqlEditorHook( hook: (props: TProps) => TResult, options?: RenderSqlEditorHookOptions ) { const inMemoryEditor = options?.inMemoryEditor ?? createInMemoryEditor() const aiAssistantState = options?.aiAssistantState ?? createAiAssistantState() const databaseSelectorState = options?.databaseSelectorState ?? createDatabaseSelectorState() const roleImpersonationState = options?.roleImpersonationState ?? createRoleImpersonationState('default', { current: async () => ({}) }) const flags = options?.flags const wrapper = ({ children }: { children: ReactNode }) => { const tree = ( {children} ) return ( {flags ? ( {tree} ) : ( tree )} ) } const result = renderHook(hook, { initialProps: options?.initialProps, wrapper, } as RenderHookOptions) return { ...result, inMemoryEditor, aiAssistantState, databaseSelectorState, roleImpersonationState, } }