Fix row.spec.ts

packages/backend-core/tests/core/utilities/mocks/licenses.ts

@@ -1,5 +1,12 @@
-import { Feature, License, Quotas } from "@budibase/types"
+import {
+  Feature,
+  License,
+  MonthlyQuotaName,
+  QuotaType,
+  QuotaUsageType,
+} from "@budibase/types"
 import cloneDeep from "lodash/cloneDeep"
+import merge from "lodash/merge"
 
 let CLOUD_FREE_LICENSE: License
 let UNLIMITED_LICENSE: License
@@ -27,18 +34,19 @@ export function initInternal(opts: {
 
 export interface UseLicenseOpts {
   features?: Feature[]
-  quotas?: Quotas
+  monthlyQuotas?: [MonthlyQuotaName, number][]
 }
 
 // LICENSES
 
 export const useLicense = (license: License, opts?: UseLicenseOpts) => {
-  if (opts) {
-    if (opts.features) {
-      license.features.push(...opts.features)
-    }
-    if (opts.quotas) {
-      license.quotas = opts.quotas
+  if (opts?.features) {
+    license.features.push(...opts.features)
+  }
+  if (opts?.monthlyQuotas) {
+    for (const [name, value] of opts.monthlyQuotas) {
+      license.quotas[QuotaType.USAGE][QuotaUsageType.MONTHLY][name].value =
+        value
     }
   }
 
@@ -57,12 +65,9 @@ export const useCloudFree = () => {
 
 // FEATURES
 
-const useFeature = (feature: Feature) => {
+const useFeature = (feature: Feature, extra?: Partial<UseLicenseOpts>) => {
   const license = cloneDeep(getCachedLicense() || UNLIMITED_LICENSE)
-  const opts: UseLicenseOpts = {
-    features: [feature],
-  }
-
+  const opts: UseLicenseOpts = merge({ features: [feature] }, extra)
   return useLicense(license, opts)
 }
 
@@ -102,8 +107,12 @@ export const useAppBuilders = () => {
   return useFeature(Feature.APP_BUILDERS)
 }
 
-export const useBudibaseAI = () => {
-  return useFeature(Feature.BUDIBASE_AI)
+export const useBudibaseAI = (opts?: { monthlyQuota?: number }) => {
+  return useFeature(Feature.BUDIBASE_AI, {
+    monthlyQuotas: [
+      [MonthlyQuotaName.BUDIBASE_AI_CREDITS, opts?.monthlyQuota || 1000],
+    ],
+  })
 }
 
 export const useAICustomConfigs = () => {

packages/server/src/api/routes/tests/row.spec.ts

@@ -39,6 +39,7 @@ import {
   RowExportFormat,
   RelationSchemaField,
   FormulaResponseType,
+  MonthlyQuotaName,
 } from "@budibase/types"
 import { generator, mocks } from "@budibase/backend-core/tests"
 import _, { merge } from "lodash"
@@ -169,6 +170,18 @@ if (descriptions.length) {
         )
       }
 
+      const resetRowUsage = async () => {
+        await config.doInContext(
+          undefined,
+          async () =>
+            await quotas.setUsage(
+              0,
+              StaticQuotaName.ROWS,
+              QuotaUsageType.STATIC
+            )
+        )
+      }
+
       const getRowUsage = async () => {
         const { total } = await config.doInContext(undefined, () =>
           quotas.getCurrentUsageValues(
@@ -206,6 +219,10 @@ if (descriptions.length) {
         table = await config.api.table.save(defaultTable())
       })
 
+      beforeEach(async () => {
+        await resetRowUsage()
+      })
+
       describe("create", () => {
         it("creates a new row successfully", async () => {
           const rowUsage = await getRowUsage()
@@ -3317,6 +3334,7 @@ if (descriptions.length) {
           beforeAll(async () => {
             mocks.licenses.useBudibaseAI()
             mocks.licenses.useAICustomConfigs()
+
             envCleanup = setEnv({
               OPENAI_API_KEY: "sk-abcdefghijklmnopqrstuvwxyz1234567890abcd",
             })