refactor(server): use feature model (#9932)

packages/backend/server/src/__tests__/app/selfhost.e2e.ts

@@ -1,7 +1,6 @@
 import { mkdirSync, writeFileSync } from 'node:fs';
 import path from 'node:path';
 
-import type { INestApplication } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
 import type { TestFn } from 'ava';
 import ava from 'ava';
@@ -10,10 +9,10 @@ import request from 'supertest';
 import { buildAppModule } from '../../app.module';
 import { Config } from '../../base';
 import { ServerService } from '../../core/config';
-import { createTestingApp, initTestingDB } from '../utils';
+import { createTestingApp, type TestingApp } from '../utils';
 
 const test = ava as TestFn<{
-  app: INestApplication;
+  app: TestingApp;
   db: PrismaClient;
 }>;
 
@@ -54,7 +53,7 @@ test.before('init selfhost server', async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.db);
+  await t.context.app.initTestingDB();
   const server = t.context.app.get(ServerService);
   // @ts-expect-error disable cache
   server._initialized = false;
@@ -188,7 +187,8 @@ test('should redirect to admin if initialized', async t => {
   t.is(res.header.location, '/admin');
 });
 
-test('should return mobile assets if visited by mobile', async t => {
+// TODO(@forehalo): return mobile when it's ready
+test.skip('should return web assets if visited by mobile', async t => {
   await t.context.db.user.create({
     data: {
       name: 'test',

packages/backend/server/src/__tests__/auth/service.spec.ts

@@ -1,4 +1,3 @@
-import { TestingModule } from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 
@@ -8,7 +7,7 @@ import { FeatureModule } from '../../core/features';
 import { QuotaModule } from '../../core/quota';
 import { UserModule } from '../../core/user';
 import { Models } from '../../models';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 const test = ava as TestFn<{
   auth: AuthService;
@@ -31,7 +30,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.db);
+  await t.context.m.initTestingDB();
   t.context.u1 = await t.context.auth.signUp('u1@affine.pro', '1');
 });
 

packages/backend/server/src/__tests__/copilot-provider.spec.ts

@@ -1,6 +1,5 @@
 /// <reference types="../global.d.ts" />
 
-import { TestingModule } from '@nestjs/testing';
 import type { ExecutionContext, TestFn } from 'ava';
 import ava from 'ava';
 
@@ -27,7 +26,7 @@ import {
   CopilotCheckHtmlExecutor,
   CopilotCheckJsonExecutor,
 } from '../plugins/copilot/workflow/executor';
-import { createTestingModule } from './utils';
+import { createTestingModule, TestingModule } from './utils';
 import { TestAssets } from './utils/copilot';
 
 type Tester = {

packages/backend/server/src/__tests__/copilot.e2e.ts

@@ -2,7 +2,6 @@
 
 import { randomUUID } from 'node:crypto';
 
-import { INestApplication } from '@nestjs/common';
 import type { TestFn } from 'ava';
 import ava from 'ava';
 import Sinon from 'sinon';
@@ -27,6 +26,7 @@ import {
   createWorkspace,
   inviteUser,
   signUp,
+  TestingApp,
 } from './utils';
 import {
   array2sse,
@@ -47,7 +47,7 @@ import {
 
 const test = ava as TestFn<{
   auth: AuthService;
-  app: INestApplication;
+  app: TestingApp;
   prompt: PromptService;
   provider: CopilotProviderService;
   storage: CopilotStorage;

packages/backend/server/src/__tests__/copilot.spec.ts

@@ -1,6 +1,3 @@
-/// <reference types="../global.d.ts" />
-
-import { TestingModule } from '@nestjs/testing';
 import type { TestFn } from 'ava';
 import ava from 'ava';
 import Sinon from 'sinon';
@@ -41,7 +38,7 @@ import {
 } from '../plugins/copilot/workflow/executor';
 import { AutoRegisteredWorkflowExecutor } from '../plugins/copilot/workflow/executor/utils';
 import { WorkflowGraphList } from '../plugins/copilot/workflow/graph';
-import { createTestingModule } from './utils';
+import { createTestingModule, TestingModule } from './utils';
 import { MockCopilotTestProvider, WorkflowTestCases } from './utils/copilot';
 
 const test = ava as TestFn<{

packages/backend/server/src/__tests__/doc/history.spec.ts

@@ -1,4 +1,3 @@
-import { TestingModule } from '@nestjs/testing';
 import type { Snapshot } from '@prisma/client';
 import { PrismaClient } from '@prisma/client';
 import test from 'ava';
@@ -7,7 +6,7 @@ import * as Sinon from 'sinon';
 import { DocStorageModule, PgWorkspaceDocStorageAdapter } from '../../core/doc';
 import { DocStorageOptions } from '../../core/doc/options';
 import { DocRecord } from '../../core/doc/storage';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 let m: TestingModule;
 let adapter: PgWorkspaceDocStorageAdapter;
@@ -24,7 +23,7 @@ test.before(async () => {
 });
 
 test.beforeEach(async () => {
-  await initTestingDB(db);
+  await m.initTestingDB();
   const options = m.get(DocStorageOptions);
   Sinon.stub(options, 'historyMaxAge').resolves(1000);
 });

packages/backend/server/src/__tests__/doc/renderer.spec.ts

@@ -75,7 +75,8 @@ test('should render correct html', async t => {
   );
 });
 
-test('should render correct mobile html', async t => {
+// TODO(@forehalo): enable it when mobile version is ready
+test.skip('should render correct mobile html', async t => {
   const res = await request(t.context.app.getHttpServer())
     .get('/workspace/xxxx/xxx')
     .set('user-agent', mobileUAString)

packages/backend/server/src/__tests__/doc/workspace.spec.ts

@@ -1,4 +1,3 @@
-import { TestingModule } from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import test from 'ava';
 import * as Sinon from 'sinon';
@@ -9,7 +8,7 @@ import {
   DocStorageModule,
   PgWorkspaceDocStorageAdapter as Adapter,
 } from '../../core/doc';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 let m: TestingModule;
 let db: PrismaClient;
@@ -35,7 +34,7 @@ test.before('init testing module', async () => {
 });
 
 test.beforeEach(async () => {
-  await initTestingDB(db);
+  await m.initTestingDB();
 });
 
 test.after.always(async () => {

packages/backend/server/src/__tests__/feature.spec.ts

@@ -1,176 +0,0 @@
-/// <reference types="../global.d.ts" />
-
-import { INestApplication } from '@nestjs/common';
-import type { TestFn } from 'ava';
-import ava from 'ava';
-
-import { Runtime } from '../base';
-import { AuthService } from '../core/auth/service';
-import {
-  FeatureManagementService,
-  FeatureModule,
-  FeatureService,
-  FeatureType,
-} from '../core/features';
-import { WorkspaceResolver } from '../core/workspaces/resolvers';
-import { createTestingApp } from './utils';
-import { WorkspaceResolverMock } from './utils/feature';
-
-const test = ava as TestFn<{
-  auth: AuthService;
-  feature: FeatureService;
-  workspace: WorkspaceResolver;
-  management: FeatureManagementService;
-  app: INestApplication;
-}>;
-
-test.beforeEach(async t => {
-  const { app } = await createTestingApp({
-    imports: [FeatureModule],
-    providers: [WorkspaceResolver],
-    tapModule: module => {
-      module
-        .overrideProvider(WorkspaceResolver)
-        .useClass(WorkspaceResolverMock);
-    },
-  });
-
-  const runtime = app.get(Runtime);
-  await runtime.set('flags/earlyAccessControl', true);
-  t.context.app = app;
-  t.context.auth = app.get(AuthService);
-  t.context.feature = app.get(FeatureService);
-  t.context.workspace = app.get(WorkspaceResolver);
-  t.context.management = app.get(FeatureManagementService);
-});
-
-test.afterEach.always(async t => {
-  await t.context.app.close();
-});
-
-test('should be able to set user feature', async t => {
-  const { auth, feature } = t.context;
-
-  const u1 = await auth.signUp('test@test.com', '123456');
-
-  const f1 = await feature.getUserFeatures(u1.id);
-  t.is(f1.length, 0, 'should be empty');
-
-  await feature.addUserFeature(u1.id, FeatureType.EarlyAccess, 'test');
-
-  const f2 = await feature.getUserFeatures(u1.id);
-  t.is(f2.length, 1, 'should have 1 feature');
-  t.is(f2[0].feature.name, FeatureType.EarlyAccess, 'should be early access');
-});
-
-test('should be able to check early access', async t => {
-  const { auth, feature, management } = t.context;
-  const u1 = await auth.signUp('test@test.com', '123456');
-
-  const f1 = await management.canEarlyAccess(u1.email);
-  t.false(f1, 'should not have early access');
-
-  await management.addEarlyAccess(u1.id);
-  const f2 = await management.canEarlyAccess(u1.email);
-  t.true(f2, 'should have early access');
-
-  const f3 = await feature.listUsersByFeature(FeatureType.EarlyAccess);
-  t.is(f3.length, 1, 'should have 1 user');
-  t.is(f3[0].id, u1.id, 'should be the same user');
-});
-
-test('should be able revert user feature', async t => {
-  const { auth, feature, management } = t.context;
-  const u1 = await auth.signUp('test@test.com', '123456');
-
-  const f1 = await management.canEarlyAccess(u1.email);
-  t.false(f1, 'should not have early access');
-
-  await management.addEarlyAccess(u1.id);
-  const f2 = await management.canEarlyAccess(u1.email);
-  t.true(f2, 'should have early access');
-  const q1 = await management.listEarlyAccess();
-  t.is(q1.length, 1, 'should have 1 user');
-  t.is(q1[0].id, u1.id, 'should be the same user');
-
-  await management.removeEarlyAccess(u1.id);
-  const f3 = await management.canEarlyAccess(u1.email);
-  t.false(f3, 'should not have early access');
-  const q2 = await management.listEarlyAccess();
-  t.is(q2.length, 0, 'should have no user');
-
-  const q3 = await feature.getUserFeatures(u1.id);
-  t.is(q3.length, 1, 'should have 1 feature');
-  t.is(q3[0].feature.name, FeatureType.EarlyAccess, 'should be early access');
-  t.is(q3[0].activated, false, 'should be deactivated');
-});
-
-test('should be same instance after reset the user feature', async t => {
-  const { auth, feature, management } = t.context;
-  const u1 = await auth.signUp('test@test.com', '123456');
-
-  await management.addEarlyAccess(u1.id);
-  const f1 = (await feature.getUserFeatures(u1.id))[0];
-
-  await management.removeEarlyAccess(u1.id);
-
-  await management.addEarlyAccess(u1.id);
-  const f2 = (await feature.getUserFeatures(u1.id))[1];
-
-  t.is(f1.feature, f2.feature, 'should be same instance');
-});
-
-test('should be able to set workspace feature', async t => {
-  const { auth, feature, workspace } = t.context;
-
-  const u1 = await auth.signUp('test@test.com', '123456');
-  const w1 = await workspace.createWorkspace(u1, null);
-
-  const f1 = await feature.getWorkspaceFeatures(w1.id);
-  t.is(f1.length, 0, 'should be empty');
-
-  await feature.addWorkspaceFeature(w1.id, FeatureType.Copilot, 'test');
-
-  const f2 = await feature.getWorkspaceFeatures(w1.id);
-  t.is(f2.length, 1, 'should have 1 feature');
-  t.is(f2[0].feature.name, FeatureType.Copilot, 'should be copilot');
-});
-
-test('should be able to check workspace feature', async t => {
-  const { auth, feature, workspace, management } = t.context;
-  const u1 = await auth.signUp('test@test.com', '123456');
-  const w1 = await workspace.createWorkspace(u1, null);
-
-  const f1 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot);
-  t.false(f1, 'should not have copilot');
-
-  await management.addWorkspaceFeatures(w1.id, FeatureType.Copilot, 'test');
-  const f2 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot);
-  t.true(f2, 'should have copilot');
-
-  const f3 = await feature.listWorkspacesByFeature(FeatureType.Copilot);
-  t.is(f3.length, 1, 'should have 1 workspace');
-  t.is(f3[0].id, w1.id, 'should be the same workspace');
-});
-
-test('should be able revert workspace feature', async t => {
-  const { auth, feature, workspace, management } = t.context;
-  const u1 = await auth.signUp('test@test.com', '123456');
-  const w1 = await workspace.createWorkspace(u1, null);
-
-  const f1 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot);
-  t.false(f1, 'should not have feature');
-
-  await management.addWorkspaceFeatures(w1.id, FeatureType.Copilot, 'test');
-  const f2 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot);
-  t.true(f2, 'should have feature');
-
-  await management.removeWorkspaceFeature(w1.id, FeatureType.Copilot);
-  const f3 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot);
-  t.false(f3, 'should not have feature');
-
-  const q3 = await feature.getWorkspaceFeatures(w1.id);
-  t.is(q3.length, 1, 'should have 1 feature');
-  t.is(q3[0].feature.name, FeatureType.Copilot, 'should be copilot');
-  t.is(q3[0].activated, false, 'should be deactivated');
-});

packages/backend/server/src/__tests__/mailer.spec.ts

@@ -5,7 +5,6 @@ import Sinon from 'sinon';
 
 import { AppModule } from '../app.module';
 import { MailService } from '../base/mailer';
-import { FeatureManagementService } from '../core/features';
 import { createTestingApp, createWorkspace, inviteUser, signUp } from './utils';
 const test = ava as TestFn<{
   app: INestApplication;
@@ -16,13 +15,6 @@ import * as renderers from '../mails';
 test.beforeEach(async t => {
   const { module, app } = await createTestingApp({
     imports: [AppModule],
-    tapModule: module => {
-      module.overrideProvider(FeatureManagementService).useValue({
-        hasWorkspaceFeature() {
-          return false;
-        },
-      });
-    },
   });
 
   const mail = module.get(MailService);

packages/backend/server/src/__tests__/models/__snapshots__/feature-user.spec.ts.md

@@ -0,0 +1,44 @@
+# Snapshot report for `src/__tests__/models/feature-user.spec.ts`
+
+The actual snapshot is saved in `feature-user.spec.ts.snap`.
+
+Generated by [AVA](https://avajs.dev).
+
+## should get user quota
+
+> free plan
+
+    {
+      blobLimit: 10485760,
+      businessBlobLimit: 104857600,
+      copilotActionLimit: 10,
+      historyPeriod: 604800000,
+      memberLimit: 3,
+      name: 'Free',
+      storageQuota: 10737418240,
+    }
+
+## should switch user quota
+
+> switch to pro plan
+
+    {
+      blobLimit: 104857600,
+      copilotActionLimit: 10,
+      historyPeriod: 2592000000,
+      memberLimit: 10,
+      name: 'Pro',
+      storageQuota: 107374182400,
+    }
+
+> switch to free plan
+
+    {
+      blobLimit: 10485760,
+      businessBlobLimit: 104857600,
+      copilotActionLimit: 10,
+      historyPeriod: 604800000,
+      memberLimit: 3,
+      name: 'Free',
+      storageQuota: 10737418240,
+    }

packages/backend/server/src/__tests__/models/__snapshots__/feature-workspace.spec.ts.md

@@ -0,0 +1,18 @@
+# Snapshot report for `src/__tests__/models/feature-workspace.spec.ts`
+
+The actual snapshot is saved in `feature-workspace.spec.ts.snap`.
+
+Generated by [AVA](https://avajs.dev).
+
+## should get workspace quota
+
+> team plan
+
+    {
+      blobLimit: 524288000,
+      historyPeriod: 2592000000,
+      memberLimit: 100,
+      name: 'Team Workspace',
+      seatQuota: 21474836480,
+      storageQuota: 2254857830400,
+    }

packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md

@@ -10,6 +10,7 @@ Generated by [AVA](https://avajs.dev).
 
     {
       blobLimit: 10485760,
+      businessBlobLimit: 104857600,
       copilotActionLimit: 10,
       historyPeriod: 604800000,
       memberLimit: 3,
@@ -23,6 +24,7 @@ Generated by [AVA](https://avajs.dev).
 
     {
       blobLimit: 10485760,
+      businessBlobLimit: 104857600,
       copilotActionLimit: 10,
       historyPeriod: 604800000,
       memberLimit: 3,

packages/backend/server/src/__tests__/models/feature-user.spec.ts

@@ -1,9 +1,8 @@
-import { TestingModule } from '@nestjs/testing';
-import { PrismaClient, User } from '@prisma/client';
+import { User } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 
-import { UserFeatureModel, UserModel } from '../../models';
-import { createTestingModule, initTestingDB } from '../utils';
+import { FeatureType, UserFeatureModel, UserModel } from '../../models';
+import { createTestingModule, TestingModule } from '../utils';
 
 interface Context {
   module: TestingModule;
@@ -21,7 +20,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.module.get(PrismaClient));
+  await t.context.module.initTestingDB();
   t.context.u1 = await t.context.module.get(UserModel).create({
     email: 'u1@affine.pro',
     registered: true,
@@ -41,7 +40,13 @@ test('should get null if user feature not found', async t => {
 test('should get user feature', async t => {
   const { model, u1 } = t.context;
   const userFeature = await model.get(u1.id, 'free_plan_v1');
-  t.is(userFeature?.feature, 'free_plan_v1');
+  t.is(userFeature?.name, 'free_plan_v1');
+});
+
+test('should get user quota', async t => {
+  const { model, u1 } = t.context;
+  const userQuota = await model.getQuota(u1.id);
+  t.snapshot(userQuota?.configs, 'free plan');
 });
 
 test('should list user features', async t => {
@@ -50,6 +55,16 @@ test('should list user features', async t => {
   t.like(await model.list(u1.id), ['free_plan_v1']);
 });
 
+test('should list user features by type', async t => {
+  const { model, u1 } = t.context;
+
+  await model.add(u1.id, 'free_plan_v1', 'test');
+  await model.add(u1.id, 'unlimited_copilot', 'test');
+
+  t.like(await model.list(u1.id, FeatureType.Quota), ['free_plan_v1']);
+  t.like(await model.list(u1.id, FeatureType.Feature), ['unlimited_copilot']);
+});
+
 test('should directly test user feature existence', async t => {
   const { model, u1 } = t.context;
 
@@ -82,14 +97,29 @@ test('should remove user feature', async t => {
   t.false((await model.list(u1.id)).includes('free_plan_v1'));
 });
 
-test('should switch user feature', async t => {
+test('should switch user quota', async t => {
   const { model, u1 } = t.context;
 
-  await model.switch(u1.id, 'free_plan_v1', 'pro_plan_v1', 'test');
+  await model.switchQuota(u1.id, 'pro_plan_v1', 'test');
+  const quota = await model.getQuota(u1.id);
+  t.snapshot(quota?.configs, 'switch to pro plan');
 
-  t.false(await model.has(u1.id, 'free_plan_v1'));
-  t.true(await model.has(u1.id, 'pro_plan_v1'));
-
-  t.false((await model.list(u1.id)).includes('free_plan_v1'));
-  t.true((await model.list(u1.id)).includes('pro_plan_v1'));
+  await model.switchQuota(u1.id, 'free_plan_v1', 'test');
+  const quota2 = await model.getQuota(u1.id);
+  t.snapshot(quota2?.configs, 'switch to free plan');
+});
+
+test('should not switch user quota if the new quota is the same as the current one', async t => {
+  const { model, u1 } = t.context;
+
+  await model.switchQuota(u1.id, 'free_plan_v1', 'test not switch');
+
+  // @ts-expect-error private
+  const quota = await model.db.userFeature.findFirst({
+    where: {
+      userId: u1.id,
+    },
+  });
+
+  t.not(quota?.reason, 'test not switch');
 });

packages/backend/server/src/__tests__/models/feature-workspace.spec.ts

@@ -1,9 +1,13 @@
-import { TestingModule } from '@nestjs/testing';
-import { PrismaClient, Workspace } from '@prisma/client';
+import { Workspace } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 
-import { UserModel, WorkspaceFeatureModel, WorkspaceModel } from '../../models';
-import { createTestingModule, initTestingDB } from '../utils';
+import {
+  FeatureType,
+  UserModel,
+  WorkspaceFeatureModel,
+  WorkspaceModel,
+} from '../../models';
+import { createTestingModule, type TestingModule } from '../utils';
 
 interface Context {
   module: TestingModule;
@@ -21,7 +25,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.module.get(PrismaClient));
+  await t.context.module.initTestingDB();
   const u1 = await t.context.module.get(UserModel).create({
     email: 'u1@affine.pro',
     registered: true,
@@ -46,18 +50,52 @@ test('should directly test workspace feature existence', async t => {
   t.false(await model.has(ws.id, 'unlimited_workspace'));
 });
 
+test('should get workspace quota', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'team_plan_v1', 'test', {
+    memberLimit: 100,
+  });
+
+  const quota = await model.getQuota(ws.id);
+  t.snapshot(quota?.configs, 'team plan');
+});
+
+test('should return null if quota removed', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'team_plan_v1', 'test', {
+    memberLimit: 100,
+  });
+
+  await model.remove(ws.id, 'team_plan_v1');
+
+  const quota = await model.getQuota(ws.id);
+  t.is(quota, null);
+});
+
 test('should list empty workspace features', async t => {
   const { model, ws } = t.context;
 
   t.deepEqual(await model.list(ws.id), []);
 });
 
+test('should list workspace features by type', async t => {
+  const { model, ws } = t.context;
+
+  await model.add(ws.id, 'unlimited_workspace', 'test');
+  await model.add(ws.id, 'team_plan_v1', 'test');
+
+  t.like(await model.list(ws.id, FeatureType.Quota), ['team_plan_v1']);
+  t.like(await model.list(ws.id, FeatureType.Feature), ['unlimited_workspace']);
+});
+
 test('should add workspace feature', async t => {
   const { model, ws } = t.context;
 
   await model.add(ws.id, 'unlimited_workspace', 'test');
   t.is(
-    (await model.get(ws.id, 'unlimited_workspace'))?.feature,
+    (await model.get(ws.id, 'unlimited_workspace'))?.name,
     'unlimited_workspace'
   );
   t.true(await model.has(ws.id, 'unlimited_workspace'));
@@ -103,25 +141,3 @@ test('should remove workspace feature', async t => {
   t.false(await model.has(ws.id, 'team_plan_v1'));
   t.false((await model.list(ws.id)).includes('team_plan_v1'));
 });
-
-test('should switch workspace feature', async t => {
-  const { model, ws } = t.context;
-
-  await model.switch(ws.id, 'team_plan_v1', 'unlimited_workspace', 'test');
-
-  t.false(await model.has(ws.id, 'team_plan_v1'));
-  t.true(await model.has(ws.id, 'unlimited_workspace'));
-
-  t.false((await model.list(ws.id)).includes('team_plan_v1'));
-  t.true((await model.list(ws.id)).includes('unlimited_workspace'));
-});
-
-test('should switch workspace feature with overrides', async t => {
-  const { model, ws } = t.context;
-
-  await model.add(ws.id, 'unlimited_workspace', 'test');
-  await model.add(ws.id, 'team_plan_v1', 'test', { memberLimit: 100 });
-  const f2 = await model.get(ws.id, 'team_plan_v1');
-
-  t.is(f2!.configs.memberLimit, 100);
-});

packages/backend/server/src/__tests__/models/feature.spec.ts

@@ -1,9 +1,8 @@
-import { TestingModule } from '@nestjs/testing';
-import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 
+import { FeatureType } from '../../models';
 import { FeatureModel } from '../../models/feature';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 interface Context {
   module: TestingModule;
@@ -20,7 +19,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.module.get(PrismaClient));
+  await t.context.module.initTestingDB();
 });
 
 test.after(async t => {
@@ -91,7 +90,12 @@ test('should get feature if extra fields exist in feature config', async t => {
 test('should create feature', async t => {
   const { feature } = t.context;
 
-  const newFeature = await feature.upsert('new_feature' as any, {});
+  const newFeature = await feature.upsert(
+    'new_feature' as any,
+    {},
+    FeatureType.Feature,
+    1
+  );
 
   t.deepEqual(newFeature.configs, {});
 });
@@ -100,10 +104,15 @@ test('should update feature', async t => {
   const { feature } = t.context;
   const freePlanFeature = await feature.get('free_plan_v1');
 
-  const newFreePlanFeature = await feature.upsert('free_plan_v1', {
-    ...freePlanFeature.configs,
-    memberLimit: 10,
-  });
+  const newFreePlanFeature = await feature.upsert(
+    'free_plan_v1',
+    {
+      ...freePlanFeature.configs,
+      memberLimit: 10,
+    },
+    FeatureType.Quota,
+    1
+  );
 
   t.deepEqual(newFreePlanFeature.configs, {
     ...freePlanFeature.configs,
@@ -113,7 +122,10 @@ test('should update feature', async t => {
 
 test('should throw if feature config is invalid when updating', async t => {
   const { feature } = t.context;
-  await t.throwsAsync(feature.upsert('free_plan_v1', {} as any), {
-    message: 'Invalid feature config for free_plan_v1',
-  });
+  await t.throwsAsync(
+    feature.upsert('free_plan_v1', {} as any, FeatureType.Quota, 1),
+    {
+      message: 'Invalid feature config for free_plan_v1',
+    }
+  );
 });

packages/backend/server/src/__tests__/models/page.spec.ts

@@ -1,4 +1,3 @@
-import { TestingModule } from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 
@@ -8,7 +7,7 @@ import { PublicPageMode } from '../../models/common';
 import { PageModel } from '../../models/page';
 import { type User, UserModel } from '../../models/user';
 import { type Workspace, WorkspaceModel } from '../../models/workspace';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 interface Context {
   config: Config;
@@ -36,7 +35,7 @@ let user: User;
 let workspace: Workspace;
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.db);
+  await t.context.module.initTestingDB();
   user = await t.context.user.create({
     email: 'test@affine.pro',
   });

packages/backend/server/src/__tests__/models/session.spec.ts

@@ -1,11 +1,10 @@
-import { TestingModule } from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 
 import { Config } from '../../base/config';
 import { SessionModel } from '../../models/session';
 import { UserModel } from '../../models/user';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 interface Context {
   config: Config;
@@ -28,7 +27,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.db);
+  await t.context.module.initTestingDB();
 });
 
 test.after(async t => {

packages/backend/server/src/__tests__/models/user.spec.ts

@@ -1,4 +1,3 @@
-import { TestingModule } from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 import Sinon from 'sinon';
@@ -7,7 +6,7 @@ import { EmailAlreadyUsed, EventBus } from '../../base';
 import { WorkspaceRole } from '../../core/permission';
 import { UserModel } from '../../models/user';
 import { WorkspaceMemberStatus } from '../../models/workspace';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 interface Context {
   module: TestingModule;
@@ -24,7 +23,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.module.get(PrismaClient));
+  await t.context.module.initTestingDB();
 });
 
 test.after(async t => {

packages/backend/server/src/__tests__/models/verification-token.spec.ts

@@ -1,4 +1,3 @@
-import { TestingModule } from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 
@@ -6,7 +5,7 @@ import {
   TokenType,
   VerificationTokenModel,
 } from '../../models/verification-token';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 interface Context {
   module: TestingModule;
@@ -25,7 +24,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.db);
+  await t.context.module.initTestingDB();
 });
 
 test.after(async t => {

packages/backend/server/src/__tests__/models/workspace.spec.ts

@@ -1,4 +1,3 @@
-import { TestingModule } from '@nestjs/testing';
 import { PrismaClient, WorkspaceMemberStatus } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 import Sinon from 'sinon';
@@ -7,7 +6,7 @@ import { Config, EventBus } from '../../base';
 import { WorkspaceRole } from '../../core/permission';
 import { UserModel } from '../../models/user';
 import { WorkspaceModel } from '../../models/workspace';
-import { createTestingModule, initTestingDB } from '../utils';
+import { createTestingModule, type TestingModule } from '../utils';
 
 interface Context {
   config: Config;
@@ -29,7 +28,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.db);
+  await t.context.module.initTestingDB();
 });
 
 test.after(async t => {

packages/backend/server/src/__tests__/nestjs/throttler.spec.ts

@@ -1,13 +1,6 @@
 import '../../plugins/config';
 
-import {
-  Controller,
-  Get,
-  HttpStatus,
-  INestApplication,
-  UseGuards,
-} from '@nestjs/common';
-import { PrismaClient } from '@prisma/client';
+import { Controller, Get, HttpStatus, UseGuards } from '@nestjs/common';
 import ava, { TestFn } from 'ava';
 import Sinon from 'sinon';
 import request, { type Response } from 'supertest';
@@ -21,12 +14,12 @@ import {
   ThrottlerStorage,
 } from '../../base/throttler';
 import { AuthService, Public } from '../../core/auth';
-import { createTestingApp, initTestingDB, internalSignIn } from '../utils';
+import { createTestingApp, internalSignIn, TestingApp } from '../utils';
 
 const test = ava as TestFn<{
   storage: ThrottlerStorage;
   cookie: string;
-  app: INestApplication;
+  app: TestingApp;
 }>;
 
 @UseGuards(CloudThrottlerGuard)
@@ -115,7 +108,7 @@ test.before(async t => {
 });
 
 test.beforeEach(async t => {
-  await initTestingDB(t.context.app.get(PrismaClient));
+  await t.context.app.initTestingDB();
   const { app } = t.context;
   const auth = app.get(AuthService);
   const u1 = await auth.signUp('u1@affine.pro', 'test');

packages/backend/server/src/__tests__/oauth/controller.spec.ts

@@ -1,6 +1,6 @@
 import '../../plugins/config';
 
-import { HttpStatus, INestApplication } from '@nestjs/common';
+import { HttpStatus } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 import Sinon from 'sinon';
@@ -15,7 +15,7 @@ import { Models } from '../../models';
 import { OAuthProviderName } from '../../plugins/oauth/config';
 import { GoogleOAuthProvider } from '../../plugins/oauth/providers/google';
 import { OAuthService } from '../../plugins/oauth/service';
-import { createTestingApp, getSession, initTestingDB } from '../utils';
+import { createTestingApp, getSession, TestingApp } from '../utils';
 
 const test = ava as TestFn<{
   auth: AuthService;
@@ -23,7 +23,7 @@ const test = ava as TestFn<{
   models: Models;
   u1: CurrentUser;
   db: PrismaClient;
-  app: INestApplication;
+  app: TestingApp;
 }>;
 
 test.before(async t => {
@@ -54,7 +54,7 @@ test.before(async t => {
 
 test.beforeEach(async t => {
   Sinon.restore();
-  await initTestingDB(t.context.db);
+  await t.context.app.initTestingDB();
   t.context.u1 = await t.context.auth.signUp('u1@affine.pro', '1');
 });
 
@@ -247,7 +247,7 @@ test('should throw if provider is invalid in callback uri', async t => {
   t.pass();
 });
 
-function mockOAuthProvider(app: INestApplication, email: string) {
+function mockOAuthProvider(app: TestingApp, email: string) {
   const provider = app.get(GoogleOAuthProvider);
   const oauth = app.get(OAuthService);
 

packages/backend/server/src/__tests__/payment/service.spec.ts

@@ -1,6 +1,5 @@
 import '../../plugins/payment';
 
-import { INestApplication } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 import Sinon from 'sinon';
@@ -11,7 +10,7 @@ import { EventBus, Runtime } from '../../base';
 import { ConfigModule } from '../../base/config';
 import { CurrentUser } from '../../core/auth';
 import { AuthService } from '../../core/auth/service';
-import { EarlyAccessType, FeatureManagementService } from '../../core/features';
+import { EarlyAccessType, FeatureService } from '../../core/features';
 import { SubscriptionService } from '../../plugins/payment/service';
 import {
   CouponType,
@@ -21,7 +20,7 @@ import {
   SubscriptionStatus,
   SubscriptionVariant,
 } from '../../plugins/payment/types';
-import { createTestingApp, initTestingDB } from '../utils';
+import { createTestingApp, type TestingApp } from '../utils';
 
 const PRO_MONTHLY = `${SubscriptionPlan.Pro}_${SubscriptionRecurring.Monthly}`;
 const PRO_YEARLY = `${SubscriptionPlan.Pro}_${SubscriptionRecurring.Yearly}`;
@@ -156,10 +155,10 @@ const sub: Stripe.Subscription = {
 const test = ava as TestFn<{
   u1: CurrentUser;
   db: PrismaClient;
-  app: INestApplication;
+  app: TestingApp;
   service: SubscriptionService;
   event: Sinon.SinonStubbedInstance<EventBus>;
-  feature: Sinon.SinonStubbedInstance<FeatureManagementService>;
+  feature: Sinon.SinonStubbedInstance<FeatureService>;
   runtime: Sinon.SinonStubbedInstance<Runtime>;
   stripe: {
     customers: Sinon.SinonStubbedInstance<Stripe.CustomersResource>;
@@ -200,8 +199,8 @@ test.before(async t => {
       AppModule,
     ],
     tapModule: m => {
-      m.overrideProvider(FeatureManagementService).useValue(
-        Sinon.createStubInstance(FeatureManagementService)
+      m.overrideProvider(FeatureService).useValue(
+        Sinon.createStubInstance(FeatureService)
       );
       m.overrideProvider(EventBus).useValue(Sinon.createStubInstance(EventBus));
       m.overrideProvider(Runtime).useValue(Sinon.createStubInstance(Runtime));
@@ -210,7 +209,7 @@ test.before(async t => {
 
   t.context.event = app.get(EventBus);
   t.context.service = app.get(SubscriptionService);
-  t.context.feature = app.get(FeatureManagementService);
+  t.context.feature = app.get(FeatureService);
   t.context.runtime = app.get(Runtime);
   t.context.db = app.get(PrismaClient);
   t.context.app = app;
@@ -232,7 +231,7 @@ test.before(async t => {
 
 test.beforeEach(async t => {
   const { db, app, stripe } = t.context;
-  await initTestingDB(db);
+  await t.context.app.initTestingDB();
   t.context.u1 = await app.get(AuthService).signUp('u1@affine.pro', '1');
 
   await db.workspace.create({

packages/backend/server/src/__tests__/quota.spec.ts

@@ -1,212 +0,0 @@
-/// <reference types="../global.d.ts" />
-
-import { TestingModule } from '@nestjs/testing';
-import type { TestFn } from 'ava';
-import ava from 'ava';
-
-import { AuthService } from '../core/auth';
-import {
-  QuotaManagementService,
-  QuotaModule,
-  QuotaService,
-  QuotaType,
-} from '../core/quota';
-import { OneGB, OneMB } from '../core/quota/constant';
-import { FreePlan, ProPlan } from '../core/quota/schema';
-import { StorageModule, WorkspaceBlobStorage } from '../core/storage';
-import { WorkspaceResolver } from '../core/workspaces/resolvers';
-import { createTestingModule } from './utils';
-import { WorkspaceResolverMock } from './utils/feature';
-
-const test = ava as TestFn<{
-  auth: AuthService;
-  quota: QuotaService;
-  quotaManager: QuotaManagementService;
-  workspace: WorkspaceResolver;
-  workspaceBlob: WorkspaceBlobStorage;
-  module: TestingModule;
-}>;
-
-test.beforeEach(async t => {
-  const module = await createTestingModule({
-    imports: [StorageModule, QuotaModule],
-    providers: [WorkspaceResolver],
-    tapModule: module => {
-      module
-        .overrideProvider(WorkspaceResolver)
-        .useClass(WorkspaceResolverMock);
-    },
-  });
-
-  t.context.module = module;
-  t.context.auth = module.get(AuthService);
-  t.context.quota = module.get(QuotaService);
-  t.context.quotaManager = module.get(QuotaManagementService);
-  t.context.workspace = module.get(WorkspaceResolver);
-  t.context.workspaceBlob = module.get(WorkspaceBlobStorage);
-});
-
-test.afterEach.always(async t => {
-  await t.context.module.close();
-});
-
-test('should be able to set quota', async t => {
-  const { auth, quota } = t.context;
-
-  const u1 = await auth.signUp('test@affine.pro', '123456');
-
-  const q1 = await quota.getUserQuota(u1.id);
-  t.truthy(q1, 'should have quota');
-  t.is(q1?.feature.name, QuotaType.FreePlanV1, 'should be free plan');
-  t.is(q1?.feature.version, 4, 'should be version 4');
-
-  await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1);
-
-  const q2 = await quota.getUserQuota(u1.id);
-  t.is(q2?.feature.name, QuotaType.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, quotaManager } = t.context;
-  const u1 = await auth.signUp('test@affine.pro', '123456');
-  const freePlan = FreePlan.configs;
-  const proPlan = ProPlan.configs;
-
-  const q1 = await quotaManager.getUserQuota(u1.id);
-  t.is(q1?.blobLimit, freePlan.blobLimit, 'should be free plan');
-  t.is(q1?.storageQuota, freePlan.storageQuota, 'should be free plan');
-
-  await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1);
-  const q2 = await quotaManager.getUserQuota(u1.id);
-  t.is(q2?.blobLimit, proPlan.blobLimit, 'should be pro plan');
-  t.is(q2?.storageQuota, proPlan.storageQuota, 'should be pro plan');
-});
-
-test('should be able revert quota', async t => {
-  const { auth, quota, quotaManager } = t.context;
-  const u1 = await auth.signUp('test@affine.pro', '123456');
-  const freePlan = FreePlan.configs;
-  const proPlan = ProPlan.configs;
-
-  const q1 = await quotaManager.getUserQuota(u1.id);
-
-  t.is(q1?.blobLimit, freePlan.blobLimit, 'should be free plan');
-  t.is(q1?.storageQuota, freePlan.storageQuota, 'should be free plan');
-
-  await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1);
-  const q2 = await quotaManager.getUserQuota(u1.id);
-  t.is(q2?.blobLimit, proPlan.blobLimit, 'should be pro plan');
-  t.is(q2?.storageQuota, proPlan.storageQuota, 'should be pro plan');
-  t.is(
-    q2?.copilotActionLimit,
-    proPlan.copilotActionLimit!,
-    'should be pro plan'
-  );
-
-  await quota.switchUserQuota(u1.id, QuotaType.FreePlanV1);
-  const q3 = await quotaManager.getUserQuota(u1.id);
-  t.is(q3?.blobLimit, freePlan.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.name, QuotaType.FreePlanV1, 'should be free plan');
-  t.is(quotas[1].feature.name, QuotaType.ProPlanV1, 'should be pro plan');
-  t.is(quotas[2].feature.name, QuotaType.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');
-});
-
-test('should be able to check quota', async t => {
-  const { auth, quotaManager } = t.context;
-  const u1 = await auth.signUp('test@affine.pro', '123456');
-  const freePlan = FreePlan.configs;
-
-  const q1 = await quotaManager.getUserQuota(u1.id);
-  t.assert(q1, 'should have quota');
-  t.is(q1.blobLimit, freePlan.blobLimit, 'should be free plan');
-  t.is(q1.storageQuota, freePlan.storageQuota, 'should be free plan');
-  t.is(q1.historyPeriod, freePlan.historyPeriod, 'should be free plan');
-  t.is(q1.memberLimit, freePlan.memberLimit, 'should be free plan');
-  t.is(
-    q1.copilotActionLimit!,
-    freePlan.copilotActionLimit!,
-    'should be free plan'
-  );
-});
-
-test('should be able to override quota', async t => {
-  const { auth, quotaManager, workspace } = t.context;
-
-  const u1 = await auth.signUp('test@affine.pro', '123456');
-  const w1 = await workspace.createWorkspace(u1, null);
-
-  const wq1 = await quotaManager.getWorkspaceUsage(w1.id);
-  t.is(wq1.blobLimit, 10 * OneMB, 'should be 10MB');
-  t.is(wq1.businessBlobLimit, 100 * OneMB, 'should be 100MB');
-  t.is(wq1.memberLimit, 3, 'should be 3');
-
-  await quotaManager.addTeamWorkspace(w1.id, 'test');
-  const wq2 = await quotaManager.getWorkspaceUsage(w1.id);
-  t.is(wq2.storageQuota, 120 * OneGB, 'should be override to 100GB');
-  t.is(wq2.businessBlobLimit, 500 * OneMB, 'should be override to 500MB');
-  t.is(wq2.memberLimit, 1, 'should be override to 1');
-
-  await quotaManager.updateWorkspaceConfig(w1.id, QuotaType.TeamPlanV1, {
-    memberLimit: 2,
-  });
-  const wq3 = await quotaManager.getWorkspaceUsage(w1.id);
-  t.is(wq3.storageQuota, 140 * OneGB, 'should be override to 120GB');
-  t.is(wq3.memberLimit, 2, 'should be override to 1');
-});
-
-test('should be able to check with workspace quota', async t => {
-  const { auth, quotaManager, workspace, workspaceBlob } = t.context;
-
-  const u1 = await auth.signUp('test@affine.pro', '123456');
-  const w1 = await workspace.createWorkspace(u1, null);
-  const w2 = await workspace.createWorkspace(u1, null);
-  const w3 = await workspace.createWorkspace(u1, null);
-  await quotaManager.addTeamWorkspace(w3.id, 'test');
-
-  {
-    const wq1 = await quotaManager.getWorkspaceUsage(w1.id);
-    t.is(wq1.usedSize, 0, 'should be 0');
-    const wq2 = await quotaManager.getWorkspaceUsage(w2.id);
-    t.is(wq2.usedSize, 0, 'should be 0');
-    const wq3 = await quotaManager.getWorkspaceUsage(w3.id);
-    t.is(wq3.usedSize, 0, 'should be 0');
-  }
-
-  {
-    await workspaceBlob.put(w1.id, 'test', Buffer.from([0, 0]));
-    await workspaceBlob.put(w2.id, 'test', Buffer.from([0, 0]));
-
-    // normal workspace
-    const wq1 = await quotaManager.getWorkspaceUsage(w1.id);
-    t.is(wq1.usedSize, 4, 'should share usage with w2');
-    const wq2 = await quotaManager.getWorkspaceUsage(w2.id);
-    t.is(wq2.usedSize, 4, 'should share usage with w1');
-
-    // workspace with quota
-    const wq3 = await quotaManager.getWorkspaceUsage(w3.id);
-    t.is(wq3.usedSize, 0, 'should not share usage with w1 and w2');
-  }
-
-  {
-    await workspaceBlob.put(w3.id, 'test', Buffer.from([0, 0, 0]));
-
-    // normal workspace
-    const wq1 = await quotaManager.getWorkspaceUsage(w1.id);
-    t.is(wq1.usedSize, 4, 'should not share usage with w3');
-    const wq2 = await quotaManager.getWorkspaceUsage(w2.id);
-    t.is(wq2.usedSize, 4, 'should not share usage with w3');
-
-    // workspace with quota
-    const wq3 = await quotaManager.getWorkspaceUsage(w3.id);
-    t.is(wq3.usedSize, 3, 'should not share usage with w1 and w2');
-  }
-});

packages/backend/server/src/__tests__/team.e2e.ts

@@ -3,7 +3,6 @@
 import { randomUUID } from 'node:crypto';
 
 import { getCurrentMailMessageCount } from '@affine-test/kit/utils/cloud';
-import { INestApplication } from '@nestjs/common';
 import { WorkspaceMemberStatus } from '@prisma/client';
 import type { TestFn } from 'ava';
 import ava from 'ava';
@@ -14,8 +13,8 @@ import { EventBus } from '../base';
 import { AuthService } from '../core/auth';
 import { DocContentService } from '../core/doc-renderer';
 import { PermissionService, WorkspaceRole } from '../core/permission';
-import { QuotaManagementService, QuotaService, QuotaType } from '../core/quota';
 import { WorkspaceType } from '../core/workspaces';
+import { Models } from '../models';
 import {
   acceptInviteById,
   approveMember,
@@ -34,19 +33,19 @@ import {
   revokeUser,
   signUp,
   sleep,
+  TestingApp,
   UserAuthedType,
 } from './utils';
 
 const test = ava as TestFn<{
-  app: INestApplication;
+  app: TestingApp;
   auth: AuthService;
   event: Sinon.SinonStubbedInstance<EventBus>;
-  quota: QuotaService;
-  quotaManager: QuotaManagementService;
+  models: Models;
   permissions: PermissionService;
 }>;
 
-test.beforeEach(async t => {
+test.before(async t => {
   const { app } = await createTestingApp({
     imports: [AppModule],
     tapModule: module => {
@@ -67,17 +66,20 @@ test.beforeEach(async t => {
   t.context.app = app;
   t.context.auth = app.get(AuthService);
   t.context.event = app.get(EventBus);
-  t.context.quota = app.get(QuotaService);
-  t.context.quotaManager = app.get(QuotaManagementService);
+  t.context.models = app.get(Models);
   t.context.permissions = app.get(PermissionService);
 });
 
-test.afterEach.always(async t => {
+test.beforeEach(async t => {
+  await t.context.app.initTestingDB();
+});
+
+test.after.always(async t => {
   await t.context.app.close();
 });
 
 const init = async (
-  app: INestApplication,
+  app: TestingApp,
   memberLimit = 10,
   prefix = randomUUID()
 ) => {
@@ -87,17 +89,15 @@ const init = async (
     `${prefix}owner@affine.pro`,
     '123456'
   );
+  const models = app.get(Models);
   {
-    const quota = app.get(QuotaService);
-    await quota.switchUserQuota(owner.id, QuotaType.ProPlanV1);
+    await models.userFeature.add(owner.id, 'pro_plan_v1', 'test');
   }
 
   const workspace = await createWorkspace(app, owner.token.token);
   const teamWorkspace = await createWorkspace(app, owner.token.token);
   {
-    const quota = app.get(QuotaManagementService);
-    await quota.addTeamWorkspace(teamWorkspace.id, 'test');
-    await quota.updateWorkspaceConfig(teamWorkspace.id, QuotaType.TeamPlanV1, {
+    models.workspaceFeature.add(teamWorkspace.id, 'team_plan_v1', 'test', {
       memberLimit,
     });
   }
@@ -264,7 +264,7 @@ test('should be able to invite multiple users', async t => {
 });
 
 test('should be able to check seat limit', async t => {
-  const { app, permissions, quotaManager } = t.context;
+  const { app, permissions, models } = t.context;
   const { invite, inviteBatch, teamWorkspace: ws } = await init(app, 4);
 
   {
@@ -274,7 +274,7 @@ test('should be able to check seat limit', async t => {
       { message: 'You have exceeded your workspace member quota.' },
       'should throw error if exceed member limit'
     );
-    await quotaManager.updateWorkspaceConfig(ws.id, QuotaType.TeamPlanV1, {
+    models.workspaceFeature.add(ws.id, 'team_plan_v1', 'test', {
       memberLimit: 5,
     });
     await t.notThrowsAsync(
@@ -323,7 +323,7 @@ test('should be able to check seat limit', async t => {
 
 test('should be able to grant team member permission', async t => {
   const { app, permissions } = t.context;
-  const { owner, teamWorkspace: ws, admin, write, read } = await init(app);
+  const { owner, teamWorkspace: ws, write, read } = await init(app);
 
   await t.throwsAsync(
     grantMember(
@@ -350,13 +350,13 @@ test('should be able to grant team member permission', async t => {
   await t.throwsAsync(
     grantMember(
       app,
-      admin.token.token,
+      write.token.token,
       ws.id,
       read.id,
       WorkspaceRole.Collaborator
     ),
     { instanceOf: Error },
-    'should throw error if not owner'
+    'should throw error if not admin'
   );
 
   {
@@ -571,7 +571,7 @@ test('should be able to approve team member', async t => {
 });
 
 test('should be able to invite by link', async t => {
-  const { app, permissions, quotaManager } = t.context;
+  const { app, permissions, models } = t.context;
   const {
     createInviteLink,
     owner,
@@ -631,7 +631,7 @@ test('should be able to invite by link', async t => {
       'should not change status'
     );
 
-    await quotaManager.updateWorkspaceConfig(tws.id, QuotaType.TeamPlanV1, {
+    models.workspaceFeature.add(tws.id, 'team_plan_v1', 'test', {
       memberLimit: 5,
     });
     await permissions.refreshSeatStatus(tws.id, 5);
@@ -646,7 +646,7 @@ test('should be able to invite by link', async t => {
       'should not change status'
     );
 
-    await quotaManager.updateWorkspaceConfig(tws.id, QuotaType.TeamPlanV1, {
+    models.workspaceFeature.add(tws.id, 'team_plan_v1', 'test', {
       memberLimit: 6,
     });
     await permissions.refreshSeatStatus(tws.id, 6);

packages/backend/server/src/__tests__/user.e2e.ts

@@ -1,20 +1,23 @@
-import type { INestApplication } from '@nestjs/common';
 import test from 'ava';
 import request from 'supertest';
 
 import { AppModule } from '../app.module';
-import { createTestingApp, currentUser, signUp } from './utils';
+import { createTestingApp, currentUser, signUp, TestingApp } from './utils';
 
-let app: INestApplication;
+let app: TestingApp;
 
-test.beforeEach(async () => {
+test.before(async () => {
   const { app: testApp } = await createTestingApp({
     imports: [AppModule],
   });
   app = testApp;
 });
 
-test.afterEach.always(async () => {
+test.beforeEach(async () => {
+  await app.initTestingDB();
+});
+
+test.after.always(async () => {
   await app.close();
 });
 

packages/backend/server/src/__tests__/utils/utils.ts

@@ -3,9 +3,13 @@ import {
   INestApplication,
   ModuleMetadata,
 } from '@nestjs/common';
-import { APP_GUARD } from '@nestjs/core';
+import { APP_GUARD, ModuleRef } from '@nestjs/core';
 import { Query, Resolver } from '@nestjs/graphql';
-import { Test, TestingModuleBuilder } from '@nestjs/testing';
+import {
+  Test,
+  TestingModule as BaseTestingModule,
+  TestingModuleBuilder,
+} from '@nestjs/testing';
 import { PrismaClient } from '@prisma/client';
 import cookieParser from 'cookie-parser';
 import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
@@ -16,7 +20,7 @@ import { AppModule, FunctionalityModules } from '../../app.module';
 import { GlobalExceptionFilter, Runtime } from '../../base';
 import { GqlModule } from '../../base/graphql';
 import { AuthGuard, AuthModule } from '../../core/auth';
-import { UserFeaturesInit1698652531198 } from '../../data/migrations/1698652531198-user-features-init';
+import { RefreshFeatures0001 } from '../../data/migrations/0001-refresh-features';
 import { ModelsModule } from '../../models';
 
 async function flushDB(client: PrismaClient) {
@@ -35,20 +39,25 @@ async function flushDB(client: PrismaClient) {
   );
 }
 
-async function initFeatureConfigs(db: PrismaClient) {
-  await UserFeaturesInit1698652531198.up(db);
-}
-
-export async function initTestingDB(db: PrismaClient) {
-  await flushDB(db);
-  await initFeatureConfigs(db);
-}
-
 interface TestingModuleMeatdata extends ModuleMetadata {
   tapModule?(m: TestingModuleBuilder): void;
   tapApp?(app: INestApplication): void;
 }
 
+const initTestingDB = async (ref: ModuleRef) => {
+  const db = ref.get(PrismaClient, { strict: false });
+  await flushDB(db);
+  await RefreshFeatures0001.up(db, ref);
+};
+
+export type TestingModule = BaseTestingModule & {
+  initTestingDB(): Promise<void>;
+};
+
+export type TestingApp = INestApplication & {
+  initTestingDB(): Promise<void>;
+};
+
 function dedupeModules(modules: NonNullable<ModuleMetadata['imports']>) {
   const map = new Map();
 
@@ -73,7 +82,7 @@ class MockResolver {
 
 export async function createTestingModule(
   moduleDef: TestingModuleMeatdata = {},
-  init = true
+  autoInitialize = true
 ) {
   // setting up
   let imports = moduleDef.imports ?? [];
@@ -107,13 +116,9 @@ export async function createTestingModule(
 
   const m = await builder.compile();
 
-  const prisma = m.get(PrismaClient);
-  if (prisma instanceof PrismaClient) {
-    await initTestingDB(prisma);
-  }
-
-  if (init) {
-    await m.init();
+  const testingModule = m as TestingModule;
+  testingModule.initTestingDB = async () => {
+    await initTestingDB(m.get(ModuleRef));
     // we got a lot smoking tests try to break nestjs
     // can't tolerate the noisy logs
     // @ts-expect-error private
@@ -123,9 +128,14 @@ export async function createTestingModule(
     const runtime = m.get(Runtime);
     // by pass password min length validation
     await runtime.set('auth/password.min', 1);
+  };
+
+  if (autoInitialize) {
+    await testingModule.initTestingDB();
+    await testingModule.init();
   }
 
-  return m;
+  return testingModule;
 }
 
 export async function createTestingApp(moduleDef: TestingModuleMeatdata = {}) {
@@ -135,7 +145,7 @@ export async function createTestingApp(moduleDef: TestingModuleMeatdata = {}) {
     cors: true,
     bodyParser: true,
     rawBody: true,
-  });
+  }) as TestingApp;
   const logger = new ConsoleLogger();
 
   logger.setLogLevels(['fatal']);
@@ -155,15 +165,14 @@ export async function createTestingApp(moduleDef: TestingModuleMeatdata = {}) {
     moduleDef.tapApp(app);
   }
 
+  await m.initTestingDB();
   await app.init();
 
-  const runtime = app.get(Runtime);
-  // by pass password min length validation
-  await runtime.set('auth/password.min', 1);
+  app.initTestingDB = m.initTestingDB.bind(m);
 
   return {
     module: m,
-    app,
+    app: app,
   };
 }
 

packages/backend/server/src/__tests__/workspace-invite.e2e.ts

@@ -2,7 +2,6 @@ import {
   getCurrentMailMessageCount,
   getLatestMailMessage,
 } from '@affine-test/kit/utils/cloud';
-import type { INestApplication } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
 import type { TestFn } from 'ava';
 import ava from 'ava';
@@ -20,17 +19,18 @@ import {
   leaveWorkspace,
   revokeUser,
   signUp,
+  TestingApp,
 } from './utils';
 
 const test = ava as TestFn<{
-  app: INestApplication;
+  app: TestingApp;
   client: PrismaClient;
   auth: AuthService;
   mail: MailService;
   models: Models;
 }>;
 
-test.beforeEach(async t => {
+test.before(async t => {
   const { app } = await createTestingApp({
     imports: [AppModule],
   });
@@ -41,7 +41,11 @@ test.beforeEach(async t => {
   t.context.models = app.get(Models);
 });
 
-test.afterEach.always(async t => {
+test.beforeEach(async t => {
+  await t.context.app.initTestingDB();
+});
+
+test.after.always(async t => {
   await t.context.app.close();
 });
 
@@ -227,15 +231,12 @@ test('should support pagination for member', async t => {
 test('should limit member count correctly', async t => {
   const { app } = t.context;
   const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1');
-  for (let i = 0; i < 10; i++) {
-    const workspace = await createWorkspace(app, u1.token.token);
-    await Promise.allSettled(
-      Array.from({ length: 10 }).map(async (_, i) =>
-        inviteUser(app, u1.token.token, workspace.id, `u${i}@affine.pro`)
-      )
-    );
-
-    const ws = await getWorkspace(app, u1.token.token, workspace.id);
-    t.assert(ws.members.length <= 3, 'failed to check member list');
-  }
+  const workspace = await createWorkspace(app, u1.token.token);
+  await Promise.allSettled(
+    Array.from({ length: 10 }).map(async (_, i) =>
+      inviteUser(app, u1.token.token, workspace.id, `u${i}@affine.pro`)
+    )
+  );
+  const ws = await getWorkspace(app, u1.token.token, workspace.id);
+  t.assert(ws.members.length <= 3, 'failed to check member list');
 });

packages/backend/server/src/__tests__/workspace.e2e.ts

@@ -1,30 +1,28 @@
-import type { INestApplication } from '@nestjs/common';
 import { PrismaClient } from '@prisma/client';
 import type { TestFn } from 'ava';
 import ava from 'ava';
 import request from 'supertest';
 
 import { AppModule } from '../app.module';
-import { WorkspaceRole } from '../core/permission/types';
 import {
   acceptInviteById,
   createTestingApp,
   createWorkspace,
   getWorkspacePublicPages,
-  grantMember,
   inviteUser,
   publishPage,
   revokePublicPage,
   signUp,
+  TestingApp,
   updateWorkspace,
 } from './utils';
 
 const test = ava as TestFn<{
-  app: INestApplication;
+  app: TestingApp;
   client: PrismaClient;
 }>;
 
-test.beforeEach(async t => {
+test.before(async t => {
   const { app } = await createTestingApp({
     imports: [AppModule],
   });
@@ -33,7 +31,11 @@ test.beforeEach(async t => {
   t.context.app = app;
 });
 
-test.afterEach.always(async t => {
+test.beforeEach(async t => {
+  await t.context.app.initTestingDB();
+});
+
+test.after.always(async t => {
   await t.context.app.close();
 });
 
@@ -132,30 +134,6 @@ test('should share a page', async t => {
     'You do not have permission to access doc page2 under Space not_exists_ws.',
     'unauthorized user can share page'
   );
-
-  await acceptInviteById(
-    app,
-    workspace.id,
-    await inviteUser(app, u1.token.token, workspace.id, u2.email)
-  );
-  const msg3 = await publishPage(app, u2.token.token, workspace.id, 'page2');
-  t.is(
-    msg3,
-    `You do not have permission to access doc page2 under Space ${workspace.id}.`,
-    'WorkspaceRole and PageRole is lower than required'
-  );
-
-  await grantMember(
-    app,
-    u1.token.token,
-    workspace.id,
-    u2.id,
-    WorkspaceRole.Admin
-  );
-
-  const invited = await publishPage(app, u2.token.token, workspace.id, 'page2');
-  t.is(invited.id, 'page2', 'failed to share page');
-
   const revoke = await revokePublicPage(
     app,
     u1.token.token,
@@ -168,9 +146,7 @@ test('should share a page', async t => {
     u1.token.token,
     workspace.id
   );
-  t.is(pages2.length, 1, 'failed to get shared pages');
-  t.is(pages2[0].id, 'page2', 'failed to get shared page: page2');
-
+  t.is(pages2.length, 0, 'failed to get shared pages');
   const msg4 = await revokePublicPage(
     app,
     u1.token.token,
@@ -179,19 +155,12 @@ test('should share a page', async t => {
   );
   t.is(msg4, 'Page is not public');
 
-  const revoked = await revokePublicPage(
-    app,
-    u1.token.token,
-    workspace.id,
-    'page2'
-  );
-  t.false(revoked.public, 'failed to revoke page');
-  const page3 = await getWorkspacePublicPages(
+  const pages3 = await getWorkspacePublicPages(
     app,
     u1.token.token,
     workspace.id
   );
-  t.is(page3.length, 0, 'failed to get shared pages');
+  t.is(pages3.length, 0, 'failed to get shared pages');
 });
 
 test('should be able to get workspace doc', async t => {

packages/backend/server/src/__tests__/workspace/blobs.e2e.ts

@@ -1,10 +1,8 @@
-import type { INestApplication } from '@nestjs/common';
 import test from 'ava';
 import request from 'supertest';
 
 import { AppModule } from '../../app.module';
-import { FeatureManagementService, FeatureType } from '../../core/features';
-import { QuotaService, QuotaType } from '../../core/quota';
+import { WorkspaceFeatureModel } from '../../models';
 import {
   collectAllBlobSizes,
   createTestingApp,
@@ -13,25 +11,35 @@ import {
   listBlobs,
   setBlob,
   signUp,
+  TestingApp,
 } from '../utils';
 
 const OneMB = 1024 * 1024;
+const RESTRICTED_QUOTA = {
+  seatQuota: 0,
+  blobLimit: OneMB,
+  storageQuota: 2 * OneMB - 1,
+  historyPeriod: 1,
+  memberLimit: 1,
+};
 
-let app: INestApplication;
-let quota: QuotaService;
-let feature: FeatureManagementService;
+let app: TestingApp;
+let model: WorkspaceFeatureModel;
 
-test.beforeEach(async () => {
+test.before(async () => {
   const { app: testApp } = await createTestingApp({
     imports: [AppModule],
   });
 
   app = testApp;
-  quota = app.get(QuotaService);
-  feature = app.get(FeatureManagementService);
+  model = app.get(WorkspaceFeatureModel);
 });
 
-test.afterEach.always(async () => {
+test.beforeEach(async () => {
+  await app.initTestingDB();
+});
+
+test.after.always(async () => {
   await app.close();
 });
 
@@ -119,32 +127,23 @@ test('should reject blob exceeded limit', async t => {
   const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1');
 
   const workspace1 = await createWorkspace(app, u1.token.token);
-  await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1);
+  await model.add(workspace1.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA);
 
-  const buffer1 = Buffer.from(Array.from({ length: OneMB + 1 }, () => 0));
+  const buffer1 = Buffer.from(
+    Array.from({ length: RESTRICTED_QUOTA.blobLimit + 1 }, () => 0)
+  );
   await t.throwsAsync(setBlob(app, u1.token.token, workspace1.id, buffer1));
-
-  await quota.switchUserQuota(u1.id, QuotaType.FreePlanV1);
-
-  const buffer2 = Buffer.from(Array.from({ length: OneMB + 1 }, () => 0));
-  await t.notThrowsAsync(setBlob(app, u1.token.token, workspace1.id, buffer2));
-
-  const buffer3 = Buffer.from(Array.from({ length: 100 * OneMB + 1 }, () => 0));
-  await t.throwsAsync(setBlob(app, u1.token.token, workspace1.id, buffer3));
 });
 
 test('should reject blob exceeded quota', async t => {
   const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1');
 
   const workspace = await createWorkspace(app, u1.token.token);
-  await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1);
+  await model.add(workspace.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA);
 
   const buffer = Buffer.from(Array.from({ length: OneMB }, () => 0));
 
-  for (let i = 0; i < 10; i++) {
-    await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer));
-  }
-
+  await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer));
   await t.throwsAsync(setBlob(app, u1.token.token, workspace.id, buffer));
 });
 
@@ -152,14 +151,10 @@ test('should accept blob even storage out of quota if workspace has unlimited fe
   const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1');
 
   const workspace = await createWorkspace(app, u1.token.token);
-  await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1);
-  feature.addWorkspaceFeatures(workspace.id, FeatureType.UnlimitedWorkspace);
+  await model.add(workspace.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA);
+  await model.add(workspace.id, 'unlimited_workspace', 'test');
 
   const buffer = Buffer.from(Array.from({ length: OneMB }, () => 0));
-
-  for (let i = 0; i < 10; i++) {
-    await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer));
-  }
-
+  await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer));
   await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer));
 });