feat: integrate user usage into apis (#5075)

packages/backend/server/tests/app.e2e.ts

@@ -45,6 +45,13 @@ class FakePrisma {
       },
     };
   }
+  get newFeaturesWaitingList() {
+    return {
+      async findUnique() {
+        return null;
+      },
+    };
+  }
 }
 
 test.beforeEach(async t => {
@@ -119,6 +126,7 @@ test('should find default user', async t => {
     })
     .expect(200)
     .expect(res => {
+      console.log(res.body);
       t.is(res.body.data.user.email, 'alex.yang@example.org');
     });
 });

packages/backend/server/tests/auth.e2e.ts

@@ -9,6 +9,11 @@ import ava, { type TestFn } from 'ava';
 import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
 
 import { AppModule } from '../src/app';
+import {
+  collectMigrations,
+  RevertCommand,
+  RunCommand,
+} from '../src/data/commands/run';
 import { MailService } from '../src/modules/auth/mailer';
 import { AuthService } from '../src/modules/auth/service';
 import {
@@ -37,6 +42,7 @@ test.beforeEach(async t => {
   await client.$disconnect();
   const module = await Test.createTestingModule({
     imports: [AppModule],
+    providers: [RevertCommand, RunCommand],
   }).compile();
   const app = module.createNestApplication();
   app.use(
@@ -52,6 +58,13 @@ test.beforeEach(async t => {
   t.context.app = app;
   t.context.auth = auth;
   t.context.mail = mail;
+
+  // 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(async t => {

packages/backend/server/tests/auth.spec.ts

@@ -4,6 +4,11 @@ import { PrismaClient } from '@prisma/client';
 import test from 'ava';
 
 import { ConfigModule } from '../src/config';
+import {
+  collectMigrations,
+  RevertCommand,
+  RunCommand,
+} from '../src/data/commands/run';
 import { GqlModule } from '../src/graphql.module';
 import { AuthModule } from '../src/modules/auth';
 import { AuthResolver } from '../src/modules/auth/resolver';
@@ -40,10 +45,19 @@ test.beforeEach(async () => {
       GqlModule,
       AuthModule,
       RateLimiterModule,
+      RevertCommand,
+      RunCommand,
     ],
   }).compile();
   authService = module.get(AuthService);
   authResolver = module.get(AuthResolver);
+
+  // 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 () => {

packages/backend/server/tests/mailer.e2e.ts

@@ -11,6 +11,11 @@ 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 { GqlModule } from '../src/graphql.module';
 import { AuthModule } from '../src/modules/auth';
 import { AuthService } from '../src/modules/auth/service';
@@ -45,8 +50,16 @@ test.beforeEach(async t => {
       AuthModule,
       RateLimiterModule,
     ],
+    providers: [RevertCommand, RunCommand],
   }).compile();
   t.context.auth = t.context.module.get(AuthService);
+
+  // init features
+  const run = t.context.module.get(RunCommand);
+  const revert = t.context.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 => {

packages/backend/server/tests/quota.spec.ts

@@ -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');
+});

packages/backend/server/tests/utils.ts

@@ -324,7 +324,7 @@ async function listBlobs(
   return res.body.data.listBlobs;
 }
 
-async function collectBlobSizes(
+async function getWorkspaceBlobsSize(
   app: INestApplication,
   token: string,
   workspaceId: string
@@ -335,14 +335,14 @@ async function collectBlobSizes(
     .send({
       query: `
             query {
-              collectBlobSizes(workspaceId: "${workspaceId}") {
-                size
+              workspace(id: "${workspaceId}") {
+                blobsSize
               }
             }
           `,
     })
     .expect(200);
-  return res.body.data.collectBlobSizes.size;
+  return res.body.data.workspace.blobsSize;
 }
 
 async function collectAllBlobSizes(
@@ -566,13 +566,13 @@ export {
   changeEmail,
   checkBlobSize,
   collectAllBlobSizes,
-  collectBlobSizes,
   createWorkspace,
   currentUser,
   flushDB,
   getInviteInfo,
   getPublicWorkspace,
   getWorkspace,
+  getWorkspaceBlobsSize,
   inviteUser,
   leaveWorkspace,
   listBlobs,

packages/backend/server/tests/workspace-blobs.spec.ts

@@ -6,17 +6,24 @@ import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
 import request from 'supertest';
 
 import { AppModule } from '../src/app';
+import {
+  collectMigrations,
+  RevertCommand,
+  RunCommand,
+} from '../src/data/commands/run';
+import { QuotaService, QuotaType } from '../src/modules/quota';
 import {
   checkBlobSize,
   collectAllBlobSizes,
-  collectBlobSizes,
   createWorkspace,
+  getWorkspaceBlobsSize,
   listBlobs,
   setBlob,
   signUp,
 } from './utils';
 
 let app: INestApplication;
+let quota: QuotaService;
 
 const client = new PrismaClient();
 
@@ -33,6 +40,7 @@ test.beforeEach(async () => {
 test.beforeEach(async () => {
   const module = await Test.createTestingModule({
     imports: [AppModule],
+    providers: [RevertCommand, RunCommand],
   }).compile();
   app = module.createNestApplication();
   app.use(
@@ -41,6 +49,15 @@ test.beforeEach(async () => {
       maxFiles: 5,
     })
   );
+  quota = module.get(QuotaService);
+
+  // 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();
+
   await app.init();
 });
 
@@ -103,7 +120,7 @@ test('should calc blobs size', async t => {
   const buffer2 = Buffer.from([0, 1]);
   await setBlob(app, u1.token.token, workspace.id, buffer2);
 
-  const size = await collectBlobSizes(app, u1.token.token, workspace.id);
+  const size = await getWorkspaceBlobsSize(app, u1.token.token, workspace.id);
   t.is(size, 4, 'failed to collect blob sizes');
 });
 
@@ -143,3 +160,39 @@ test('should calc all blobs size', async t => {
   );
   t.is(size2, -1, 'failed to check blob size');
 });
+
+test('should be able calc quota after switch plan', async t => {
+  const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1');
+
+  const workspace1 = await createWorkspace(app, u1.token.token);
+
+  const buffer1 = Buffer.from([0, 0]);
+  await setBlob(app, u1.token.token, workspace1.id, buffer1);
+  const buffer2 = Buffer.from([0, 1]);
+  await setBlob(app, u1.token.token, workspace1.id, buffer2);
+
+  const workspace2 = await createWorkspace(app, u1.token.token);
+
+  const buffer3 = Buffer.from([0, 0]);
+  await setBlob(app, u1.token.token, workspace2.id, buffer3);
+  const buffer4 = Buffer.from([0, 1]);
+  await setBlob(app, u1.token.token, workspace2.id, buffer4);
+
+  const size1 = await checkBlobSize(
+    app,
+    u1.token.token,
+    workspace1.id,
+    10 * 1024 * 1024 * 1024 - 8
+  );
+  t.is(size1, 0, 'failed to check free plan blob size');
+
+  quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1);
+
+  const size2 = await checkBlobSize(
+    app,
+    u1.token.token,
+    workspace1.id,
+    100 * 1024 * 1024 * 1024 - 8
+  );
+  t.is(size2, 0, 'failed to check pro plan blob size');
+});

packages/backend/server/tests/workspace-invite.e2e.ts

@@ -9,6 +9,11 @@ import ava, { type TestFn } from 'ava';
 import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
 
 import { AppModule } from '../src/app';
+import {
+  collectMigrations,
+  RevertCommand,
+  RunCommand,
+} from '../src/data/commands/run';
 import { MailService } from '../src/modules/auth/mailer';
 import { AuthService } from '../src/modules/auth/service';
 import {
@@ -39,6 +44,7 @@ test.beforeEach(async t => {
   await client.$disconnect();
   const module = await Test.createTestingModule({
     imports: [AppModule],
+    providers: [RevertCommand, RunCommand],
   }).compile();
   const app = module.createNestApplication();
   app.use(
@@ -51,9 +57,17 @@ test.beforeEach(async t => {
 
   const auth = module.get(AuthService);
   const mail = module.get(MailService);
+
   t.context.app = app;
   t.context.auth = auth;
   t.context.mail = mail;
+
+  // 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 => {

packages/backend/server/tests/workspace-usage.spec.ts

@@ -4,6 +4,7 @@ import ava, { type TestFn } from 'ava';
 import { stub } from 'sinon';
 
 import { AppModule } from '../src/app';
+import { Quotas } from '../src/modules/quota';
 import { UsersService } from '../src/modules/users';
 import { PermissionService } from '../src/modules/workspaces/permission';
 import { WorkspaceResolver } from '../src/modules/workspaces/resolver';
@@ -20,6 +21,9 @@ class FakePermission {
       user: new FakePrisma().fakeUser,
     };
   }
+  async getOwnedWorkspaces() {
+    return [''];
+  }
 }
 
 const fakeUserService = {
@@ -42,6 +46,19 @@ test.beforeEach(async t => {
           return [];
         },
       },
+      userFeatures: {
+        async count() {
+          return 1;
+        },
+        async findFirst() {
+          return {
+            createdAt: new Date(),
+            expiredAt: new Date(),
+            reason: '',
+            feature: Quotas[0],
+          };
+        },
+      },
     })
     .overrideProvider(PermissionService)
     .useClass(FakePermission)

packages/backend/server/tests/workspace.e2e.ts

@@ -6,6 +6,11 @@ import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
 import request from 'supertest';
 
 import { AppModule } from '../src/app';
+import {
+  collectMigrations,
+  RevertCommand,
+  RunCommand,
+} from '../src/data/commands/run';
 import {
   acceptInviteById,
   createWorkspace,
@@ -34,6 +39,7 @@ test.beforeEach(async t => {
   await client.$disconnect();
   const module = await Test.createTestingModule({
     imports: [AppModule],
+    providers: [RevertCommand, RunCommand],
   }).compile();
   const app = module.createNestApplication();
   app.use(
@@ -45,6 +51,13 @@ test.beforeEach(async t => {
   await app.init();
   t.context.client = client;
   t.context.app = app;
+
+  // 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 => {