feat: integrate user usage into apis (#5075)

2026-09-25 04:56:02 +08:00 · 2023-12-14 09:50:37 +00:00
parent 63de73a815
commit ad23ead5e4
36 changed files with 984 additions and 282 deletions
@@ -0,0 +1,153 @@
+/// <reference types="../src/global.d.ts" />
+
+import { Test, TestingModule } from '@nestjs/testing';
+import { PrismaClient } from '@prisma/client';
+import ava, { type TestFn } from 'ava';
+
+import { ConfigModule } from '../src/config';
+import {
+  collectMigrations,
+  RevertCommand,
+  RunCommand,
+} from '../src/data/commands/run';
+import { AuthModule } from '../src/modules/auth';
+import { AuthService } from '../src/modules/auth/service';
+import {
+  QuotaManagementService,
+  QuotaModule,
+  Quotas,
+  QuotaService,
+  QuotaType,
+} from '../src/modules/quota';
+import { PrismaModule } from '../src/prisma';
+import { StorageModule } from '../src/storage';
+import { RateLimiterModule } from '../src/throttler';
+
+const test = ava as TestFn<{
+  auth: AuthService;
+  quota: QuotaService;
+  storageQuota: QuotaManagementService;
+  app: TestingModule;
+}>;
+
+// cleanup database before each test
+test.beforeEach(async () => {
+  const client = new PrismaClient();
+  await client.$connect();
+  await client.user.deleteMany({});
+  await client.$disconnect();
+});
+
+test.beforeEach(async t => {
+  const module = await Test.createTestingModule({
+    imports: [
+      ConfigModule.forRoot({
+        auth: {
+          accessTokenExpiresIn: 1,
+          refreshTokenExpiresIn: 1,
+          leeway: 1,
+        },
+        host: 'example.org',
+        https: true,
+      }),
+      StorageModule.forRoot(),
+      PrismaModule,
+      AuthModule,
+      QuotaModule,
+      RateLimiterModule,
+      RevertCommand,
+      RunCommand,
+    ],
+  }).compile();
+
+  const quota = module.get(QuotaService);
+  const storageQuota = module.get(QuotaManagementService);
+  const auth = module.get(AuthService);
+
+  t.context.app = module;
+  t.context.quota = quota;
+  t.context.storageQuota = storageQuota;
+  t.context.auth = auth;
+
+  // init features
+  const run = module.get(RunCommand);
+  const revert = module.get(RevertCommand);
+  const migrations = await collectMigrations();
+  await Promise.allSettled(migrations.map(m => revert.run([m.name])));
+  await run.run();
+});
+
+test.afterEach.always(async t => {
+  await t.context.app.close();
+});
+
+test('should be able to set quota', async t => {
+  const { auth, quota } = t.context;
+
+  const u1 = await auth.signUp('DarkSky', 'darksky@example.org', '123456');
+
+  const q1 = await quota.getUserQuota(u1.id);
+  t.truthy(q1, 'should have quota');
+  t.is(q1?.feature.feature, QuotaType.Quota_FreePlanV1, 'should be free plan');
+
+  await quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1);
+
+  const q2 = await quota.getUserQuota(u1.id);
+  t.is(q2?.feature.feature, QuotaType.Quota_ProPlanV1, 'should be pro plan');
+
+  const fail = quota.switchUserQuota(u1.id, 'not_exists_plan_v1' as QuotaType);
+  await t.throwsAsync(fail, { instanceOf: Error }, 'should throw error');
+});
+
+test('should be able to check storage quota', async t => {
+  const { auth, quota, storageQuota } = t.context;
+  const u1 = await auth.signUp('DarkSky', 'darksky@example.org', '123456');
+
+  const q1 = await storageQuota.getUserQuota(u1.id);
+  t.is(q1?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan');
+  t.is(q1?.storageQuota, Quotas[0].configs.storageQuota, 'should be free plan');
+
+  await quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1);
+  const q2 = await storageQuota.getUserQuota(u1.id);
+  t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan');
+  t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan');
+});
+
+test('should be able revert quota', async t => {
+  const { auth, quota, storageQuota } = t.context;
+  const u1 = await auth.signUp('DarkSky', 'darksky@example.org', '123456');
+
+  const q1 = await storageQuota.getUserQuota(u1.id);
+  t.is(q1?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan');
+  t.is(q1?.storageQuota, Quotas[0].configs.storageQuota, 'should be free plan');
+
+  await quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1);
+  const q2 = await storageQuota.getUserQuota(u1.id);
+  t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan');
+  t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan');
+
+  await quota.switchUserQuota(u1.id, QuotaType.Quota_FreePlanV1);
+  const q3 = await storageQuota.getUserQuota(u1.id);
+  t.is(q3?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan');
+
+  const quotas = await quota.getUserQuotas(u1.id);
+  t.is(quotas.length, 3, 'should have 3 quotas');
+  t.is(
+    quotas[0].feature.feature,
+    QuotaType.Quota_FreePlanV1,
+    'should be free plan'
+  );
+  t.is(
+    quotas[1].feature.feature,
+    QuotaType.Quota_ProPlanV1,
+    'should be pro plan'
+  );
+  t.is(
+    quotas[2].feature.feature,
+    QuotaType.Quota_FreePlanV1,
+    'should be free plan'
+  );
+  t.is(quotas[0].activated, false, 'should be activated');
+  t.is(quotas[1].activated, false, 'should be activated');
+  t.is(quotas[2].activated, true, 'should be activated');
+});