From 5b0ad571ecf2065d2366e6ff0e9de0599052a530 Mon Sep 17 00:00:00 2001
From: Sam Rose <hello@samwho.dev>
Date: Mon, 3 Mar 2025 16:18:25 +0000
Subject: [PATCH] De-mock AI column search tests.

---
 .../src/api/routes/tests/search.spec.ts       | 23 ++++++++-----------
 1 file changed, 10 insertions(+), 13 deletions(-)

diff --git a/packages/server/src/api/routes/tests/search.spec.ts b/packages/server/src/api/routes/tests/search.spec.ts
index caa651f3bb..e115297ee9 100644
--- a/packages/server/src/api/routes/tests/search.spec.ts
+++ b/packages/server/src/api/routes/tests/search.spec.ts
@@ -6,6 +6,7 @@ import {
   docIds,
   MAX_VALID_DATE,
   MIN_VALID_DATE,
+  setEnv,
   SQLITE_DESIGN_DOC_ID,
   utils,
   withEnv as withCoreEnv,
@@ -43,19 +44,7 @@ import { generator, structures, mocks } from "@budibase/backend-core/tests"
 import { DEFAULT_EMPLOYEE_TABLE_SCHEMA } from "../../../db/defaultData/datasource_bb_default"
 import { generateRowIdField } from "../../../integrations/utils"
 import { cloneDeep } from "lodash/fp"
-
-jest.mock("@budibase/pro", () => ({
-  ...jest.requireActual("@budibase/pro"),
-  ai: {
-    LargeLanguageModel: {
-      forCurrentTenant: async () => ({
-        llm: {},
-        run: jest.fn(() => `Mock LLM Response`),
-        buildPromptFromAIOperation: jest.fn(),
-      }),
-    },
-  },
-}))
+import { mockChatGPTResponse } from "../../../tests/utilities/mocks/openai"
 
 const descriptions = datasourceDescribe({ plus: true })
 
@@ -1896,11 +1885,15 @@ if (descriptions.length) {
               !isInMemory &&
               describe("AI Column", () => {
                 const UNEXISTING_AI_COLUMN = "Real LLM Response"
+                let envCleanup: () => void
 
                 beforeAll(async () => {
                   mocks.licenses.useBudibaseAI()
                   mocks.licenses.useAICustomConfigs()
 
+                  envCleanup = setEnv({ OPENAI_API_KEY: "mock" })
+                  mockChatGPTResponse("Mock LLM Response")
+
                   tableOrViewId = await createTableOrView({
                     product: { name: "product", type: FieldType.STRING },
                     ai: {
@@ -1917,6 +1910,10 @@ if (descriptions.length) {
                   ])
                 })
 
+                afterAll(() => {
+                  envCleanup()
+                })
+
                 describe("equal", () => {
                   it("successfully finds rows based on AI column", async () => {
                     await expectQuery({