diff --git a/packages/backend-core/tests/core/utilities/mocks/licenses.ts b/packages/backend-core/tests/core/utilities/mocks/licenses.ts index 5ba6fb36a1..436e915b81 100644 --- a/packages/backend-core/tests/core/utilities/mocks/licenses.ts +++ b/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) => { 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 = () => { diff --git a/packages/server/src/api/routes/tests/row.spec.ts b/packages/server/src/api/routes/tests/row.spec.ts index 576f0bb663..70a0bb2de9 100644 --- a/packages/server/src/api/routes/tests/row.spec.ts +++ b/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", })