feat: integrate user usage into apis (#5075)

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