From 9ef8599fb81d3317be2e8fa53690f1465eb29c3a Mon Sep 17 00:00:00 2001
From: Sam Rose <hello@samwho.dev>
Date: Thu, 10 Apr 2025 16:18:25 +0100
Subject: [PATCH] Add tests for budibase AI self host quota.

---
 packages/pro                                  |  2 +-
 .../server/src/api/routes/tests/ai.spec.ts    | 28 +++++++++++++++++--
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/packages/pro b/packages/pro
index 172d40f475..af8e1f1e53 160000
--- a/packages/pro
+++ b/packages/pro
@@ -1 +1 @@
-Subproject commit 172d40f4754b4cefc424ac63c4dfb7327cadc92d
+Subproject commit af8e1f1e53b9cb9a7d4b773e68b5f0344a895788
diff --git a/packages/server/src/api/routes/tests/ai.spec.ts b/packages/server/src/api/routes/tests/ai.spec.ts
index c983cd508d..c29b70ef16 100644
--- a/packages/server/src/api/routes/tests/ai.spec.ts
+++ b/packages/server/src/api/routes/tests/ai.spec.ts
@@ -1,7 +1,13 @@
 import { mockChatGPTResponse } from "../../../tests/utilities/mocks/ai/openai"
 import TestConfiguration from "../../../tests/utilities/TestConfiguration"
 import nock from "nock"
-import { configs, env, features, setEnv } from "@budibase/backend-core"
+import {
+  configs,
+  env,
+  features,
+  setEnv,
+  StaticDatabases,
+} from "@budibase/backend-core"
 import {
   AIInnerConfig,
   ConfigType,
@@ -10,11 +16,13 @@ import {
   PlanModel,
   PlanType,
   ProviderConfig,
+  QuotaUsage,
 } from "@budibase/types"
 import { context } from "@budibase/backend-core"
 import { mocks } from "@budibase/backend-core/tests"
 import { MockLLMResponseFn } from "../../../tests/utilities/mocks/ai"
 import { mockAnthropicResponse } from "../../../tests/utilities/mocks/ai/anthropic"
+import { quotas } from "@budibase/pro"
 
 function dedent(str: string) {
   return str
@@ -301,13 +309,29 @@ describe("BudibaseAI", () => {
         .reply(200, license)
     })
 
-    it("handles correct chat response", async () => {
+    async function getQuotaUsage() {
+      return await context.doInSelfHostTenantUsingCloud(
+        config.getTenantId(),
+        async () => {
+          return await quotas.getQuotaUsage()
+        }
+      )
+    }
+
+    it.only("handles correct chat response", async () => {
+      let usage = await getQuotaUsage()
+      expect(usage._id).toBe(`quota_usage_${config.getTenantId()}`)
+      expect(usage.monthly.current.budibaseAICredits).toBe(0)
+
       mockChatGPTResponse("Hi there!")
       const { message } = await config.api.ai.chat({
         messages: [{ role: "user", content: "Hello!" }],
         licenseKey: licenseKey,
       })
       expect(message).toBe("Hi there!")
+
+      usage = await getQuotaUsage()
+      expect(usage.monthly.current.budibaseAICredits).toBeGreaterThan(0)
     })
 
     it("handles chat response error", async () => {