From 33958be6d2eebc2cbd33944bfe4b189607825b9e Mon Sep 17 00:00:00 2001
From: Saxon Fletcher <saxonafletcher@gmail.com>
Date: Mon, 28 Jul 2025 15:44:54 +1000
Subject: [PATCH] setup tests

---
 .../Database/Tests/Tests.constants.ts         |  1 +
 .../interfaces/Database/Tests/TestsList.tsx   |  3 +--
 .../pages/project/[ref]/database/tests.tsx    | 23 ++++++++++++++++---
 3 files changed, 22 insertions(+), 5 deletions(-)
 create mode 100644 apps/studio/components/interfaces/Database/Tests/Tests.constants.ts

diff --git a/apps/studio/components/interfaces/Database/Tests/Tests.constants.ts b/apps/studio/components/interfaces/Database/Tests/Tests.constants.ts
new file mode 100644
index 00000000000..d20769bc684
--- /dev/null
+++ b/apps/studio/components/interfaces/Database/Tests/Tests.constants.ts
@@ -0,0 +1 @@
+export const SETUP_TEST_PREFIX = '00_setup'
diff --git a/apps/studio/components/interfaces/Database/Tests/TestsList.tsx b/apps/studio/components/interfaces/Database/Tests/TestsList.tsx
index 7644ed8ec15..c50976a0d7c 100644
--- a/apps/studio/components/interfaces/Database/Tests/TestsList.tsx
+++ b/apps/studio/components/interfaces/Database/Tests/TestsList.tsx
@@ -3,8 +3,7 @@ import { Alert, Card, CardHeader, CardTitle } from 'ui'
 import ShimmeringLoader from 'components/ui/ShimmeringLoader'
 import { DatabaseTest, DatabaseTestStatus } from 'data/database-tests/database-tests-query'
 import TestRow from './TestRow'
-
-const SETUP_TEST_PREFIX = '00_setup'
+import { SETUP_TEST_PREFIX } from './Tests.constants'
 
 interface DatabaseTestsListProps {
   tests: DatabaseTest[]
diff --git a/apps/studio/pages/project/[ref]/database/tests.tsx b/apps/studio/pages/project/[ref]/database/tests.tsx
index 53ecc5e039e..9422fa75d54 100644
--- a/apps/studio/pages/project/[ref]/database/tests.tsx
+++ b/apps/studio/pages/project/[ref]/database/tests.tsx
@@ -25,6 +25,7 @@ import { useDatabaseTestQuery, getDatabaseTest } from 'data/database-tests/datab
 import { useExecuteSqlMutation } from 'data/sql/execute-sql-mutation'
 import { useQueryClient } from '@tanstack/react-query'
 import { databaseTestsKeys } from 'data/database-tests/database-tests-key'
+import { SETUP_TEST_PREFIX } from 'components/interfaces/Database/Tests/Tests.constants'
 
 const DatabaseTestsPage: NextPageWithLayout = () => {
   const { setEditorPanel } = useAppStateSnapshot()
@@ -46,6 +47,9 @@ const DatabaseTestsPage: NextPageWithLayout = () => {
     projectRef: project?.ref,
   })
 
+  // Cache IDs of setup tests (begin with 00_setup)
+  const setupTestIds = (tests ?? []).filter((t) => t.name.startsWith('00_setup')).map((t) => t.id)
+
   const currentTestId = executionQueue.length > 0 ? executionQueue[0] : undefined
 
   // Lazy detail hook for the test at the head of the queue
@@ -119,7 +123,7 @@ rollback;
   const onRunAllTests = () => {
     if (!tests || tests.length === 0) return
 
-    const ids = tests.filter((t) => !t.name.startsWith('000')).map((t) => t.id)
+    const ids = tests.filter((t) => !t.name.startsWith(SETUP_TEST_PREFIX)).map((t) => t.id)
     const initialStatuses: Record<string, DatabaseTestStatus> = {}
     ids.forEach((id) => {
       initialStatuses[id] = 'queued'
@@ -146,11 +150,24 @@ rollback;
 
         if (sql.length === 0) throw new Error('Empty SQL')
 
-        // Execute (prepend setup? For simplicity skip for now)
+        // Build combined setup SQL
+        let combinedSetup = ''
+        for (const sid of setupTestIds) {
+          const sDetail = await queryClient.fetchQuery(
+            databaseTestsKeys.detail(project?.ref, sid),
+            () => getDatabaseTest({ projectRef: project?.ref, id: sid }),
+            { staleTime: 60_000 }
+          )
+          combinedSetup += (sDetail?.query ?? '') + '\n\n'
+        }
+
+        const sqlToRun = combinedSetup.length > 0 ? combinedSetup + sql : sql
+
+        // Execute
         await executeSql({
           projectRef: project.ref,
           connectionString: project.connectionString,
-          sql,
+          sql: sqlToRun,
         })
 
         setStatuses((prev) => ({ ...prev, [testId]: 'passed' }))