feat(server): team quota (#8955)

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

@@ -1,7 +1,6 @@
 /// <reference types="../src/global.d.ts" />
 
-import { INestApplication, Injectable } from '@nestjs/common';
-import { PrismaClient } from '@prisma/client';
+import { INestApplication } from '@nestjs/common';
 import type { TestFn } from 'ava';
 import ava from 'ava';
 
@@ -12,32 +11,10 @@ import {
   FeatureService,
   FeatureType,
 } from '../src/core/features';
-import { Permission } from '../src/core/permission';
-import { UserType } from '../src/core/user/types';
 import { WorkspaceResolver } from '../src/core/workspaces/resolvers';
 import { Config, ConfigModule } from '../src/fundamentals/config';
 import { createTestingApp } from './utils';
-
-@Injectable()
-class WorkspaceResolverMock {
-  constructor(private readonly prisma: PrismaClient) {}
-
-  async createWorkspace(user: UserType, _init: null) {
-    const workspace = await this.prisma.workspace.create({
-      data: {
-        public: false,
-        permissions: {
-          create: {
-            type: Permission.Owner,
-            userId: user.id,
-            accepted: true,
-          },
-        },
-      },
-    });
-    return workspace;
-  }
-}
+import { WorkspaceResolverMock } from './utils/feature';
 
 const test = ava as TestFn<{
   auth: AuthService;
@@ -105,7 +82,7 @@ test('should be able to check early access', async t => {
   const f2 = await management.canEarlyAccess(u1.email);
   t.true(f2, 'should have early access');
 
-  const f3 = await feature.listFeatureUsers(FeatureType.EarlyAccess);
+  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');
 });
@@ -179,7 +156,7 @@ test('should be able to check workspace feature', async t => {
   const f2 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot);
   t.true(f2, 'should have copilot');
 
-  const f3 = await feature.listFeatureWorkspaces(FeatureType.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');
 });

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

@@ -11,29 +11,41 @@ import {
   QuotaService,
   QuotaType,
 } from '../src/core/quota';
+import { OneGB, OneMB } from '../src/core/quota/constant';
 import { FreePlan, ProPlan } from '../src/core/quota/schema';
 import { StorageModule } from '../src/core/storage';
+import { WorkspaceResolver } from '../src/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;
   module: TestingModule;
 }>;
 
 test.beforeEach(async t => {
   const module = await createTestingModule({
     imports: [StorageModule, QuotaModule],
+    providers: [WorkspaceResolver],
+    tapModule: module => {
+      module
+        .overrideProvider(WorkspaceResolver)
+        .useClass(WorkspaceResolverMock);
+    },
   });
 
   const quota = module.get(QuotaService);
   const quotaManager = module.get(QuotaManagementService);
+  const workspace = module.get(WorkspaceResolver);
   const auth = module.get(AuthService);
 
   t.context.module = module;
   t.context.quota = quota;
   t.context.quotaManager = quotaManager;
+  t.context.workspace = workspace;
   t.context.auth = auth;
 });
 
@@ -128,3 +140,28 @@ test('should be able to check quota', async t => {
     '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');
+});

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

@@ -0,0 +1,287 @@
+/// <reference types="../src/global.d.ts" />
+
+import { INestApplication } from '@nestjs/common';
+import { WorkspaceMemberStatus } from '@prisma/client';
+import type { TestFn } from 'ava';
+import ava from 'ava';
+
+import { AppModule } from '../src/app.module';
+import { AuthService } from '../src/core/auth';
+import { Permission, PermissionService } from '../src/core/permission';
+import {
+  QuotaManagementService,
+  QuotaService,
+  QuotaType,
+} from '../src/core/quota';
+import {
+  acceptInviteById,
+  createTestingApp,
+  createWorkspace,
+  grantMember,
+  inviteLink,
+  inviteUser,
+  inviteUsers,
+  leaveWorkspace,
+  PermissionEnum,
+  signUp,
+  sleep,
+  UserAuthedType,
+} from './utils';
+
+const test = ava as TestFn<{
+  app: INestApplication;
+  auth: AuthService;
+  quota: QuotaService;
+  quotaManager: QuotaManagementService;
+  permissions: PermissionService;
+}>;
+
+test.beforeEach(async t => {
+  const { app } = await createTestingApp({
+    imports: [AppModule],
+  });
+
+  const quota = app.get(QuotaService);
+  const quotaManager = app.get(QuotaManagementService);
+  const permissions = app.get(PermissionService);
+  const auth = app.get(AuthService);
+
+  t.context.app = app;
+  t.context.quota = quota;
+  t.context.quotaManager = quotaManager;
+  t.context.permissions = permissions;
+  t.context.auth = auth;
+});
+
+test.afterEach.always(async t => {
+  await t.context.app.close();
+});
+
+const init = async (app: INestApplication, memberLimit = 10) => {
+  const owner = await signUp(app, 'test', 'test@affine.pro', '123456');
+  const ws = await createWorkspace(app, owner.token.token);
+
+  const quota = app.get(QuotaManagementService);
+  await quota.addTeamWorkspace(ws.id, 'test');
+  await quota.updateWorkspaceConfig(ws.id, QuotaType.TeamPlanV1, {
+    memberLimit,
+  });
+
+  const invite = async (
+    email: string,
+    permission: PermissionEnum = 'Write'
+  ) => {
+    const member = await signUp(app, email.split('@')[0], email, '123456');
+    const inviteId = await inviteUser(
+      app,
+      owner.token.token,
+      ws.id,
+      member.email,
+      permission
+    );
+    await acceptInviteById(app, ws.id, inviteId);
+    return member;
+  };
+
+  const inviteBatch = async (emails: string[]) => {
+    const members = [];
+    for (const email of emails) {
+      const member = await signUp(app, email.split('@')[0], email, '123456');
+      members.push(member);
+    }
+    const invites = await inviteUsers(app, owner.token.token, ws.id, emails);
+    return [members, invites] as const;
+  };
+
+  const createInviteLink = async () => {
+    const inviteId = await inviteLink(app, owner.token.token, ws.id, 'OneDay');
+    return async (email: string): Promise<UserAuthedType> => {
+      const member = await signUp(app, email.split('@')[0], email, '123456');
+      await acceptInviteById(app, ws.id, inviteId, false, member.token.token);
+      return member;
+    };
+  };
+
+  const admin = await invite('admin@affine.pro', 'Admin');
+  const write = await invite('member1@affine.pro');
+  const read = await invite('member2@affine.pro', 'Read');
+
+  return {
+    invite,
+    inviteBatch,
+    createInviteLink,
+    owner,
+    ws,
+    admin,
+    write,
+    read,
+  };
+};
+
+test('should be able to check seat limit', async t => {
+  const { app, permissions, quotaManager } = t.context;
+  const { invite, inviteBatch, ws } = await init(app, 4);
+
+  {
+    // invite
+    await t.throwsAsync(
+      invite('member3@affine.pro', 'Read'),
+      { message: 'You have exceeded your workspace member quota.' },
+      'should throw error if exceed member limit'
+    );
+    await quotaManager.updateWorkspaceConfig(ws.id, QuotaType.TeamPlanV1, {
+      memberLimit: 5,
+    });
+    await t.notThrowsAsync(
+      invite('member4@affine.pro', 'Read'),
+      'should not throw error if not exceed member limit'
+    );
+  }
+
+  {
+    const members1 = inviteBatch(['member5@affine.pro']);
+    // invite batch
+    await t.notThrowsAsync(
+      members1,
+      'should not throw error in batch invite event reach limit'
+    );
+
+    t.is(
+      await permissions.getWorkspaceMemberStatus(
+        ws.id,
+        (await members1)[0][0].id
+      ),
+      WorkspaceMemberStatus.NeedMoreSeat,
+      'should be able to check member status'
+    );
+
+    // refresh seat, fifo
+    sleep(1000);
+    const [[members2]] = await inviteBatch(['member6@affine.pro']);
+    await permissions.refreshSeatStatus(ws.id, 6);
+
+    t.is(
+      await permissions.getWorkspaceMemberStatus(
+        ws.id,
+        (await members1)[0][0].id
+      ),
+      WorkspaceMemberStatus.Accepted,
+      'should become accepted after refresh'
+    );
+    t.is(
+      await permissions.getWorkspaceMemberStatus(ws.id, members2.id),
+      WorkspaceMemberStatus.NeedMoreSeat,
+      'should not change status'
+    );
+  }
+});
+
+test('should be able to grant team member permission', async t => {
+  const { app, permissions } = t.context;
+  const { owner, ws, admin, write, read } = await init(app);
+
+  await t.throwsAsync(
+    grantMember(app, read.token.token, ws.id, write.id, 'Write'),
+    { instanceOf: Error },
+    'should throw error if not owner'
+  );
+  await t.throwsAsync(
+    grantMember(app, write.token.token, ws.id, read.id, 'Write'),
+    { instanceOf: Error },
+    'should throw error if not owner'
+  );
+  await t.throwsAsync(
+    grantMember(app, admin.token.token, ws.id, read.id, 'Write'),
+    { instanceOf: Error },
+    'should throw error if not owner'
+  );
+
+  {
+    // owner should be able to grant permission
+    t.true(
+      await permissions.tryCheckWorkspaceIs(ws.id, read.id, Permission.Read),
+      'should be able to check permission'
+    );
+    t.truthy(
+      await grantMember(app, owner.token.token, ws.id, read.id, 'Admin'),
+      'should be able to grant permission'
+    );
+    t.true(
+      await permissions.tryCheckWorkspaceIs(ws.id, read.id, Permission.Admin),
+      'should be able to check permission'
+    );
+  }
+});
+
+test('should be able to leave workspace', async t => {
+  const { app } = t.context;
+  const { owner, ws, admin, write, read } = await init(app);
+
+  t.false(
+    await leaveWorkspace(app, owner.token.token, ws.id),
+    'owner should not be able to leave workspace'
+  );
+  t.true(
+    await leaveWorkspace(app, admin.token.token, ws.id),
+    'admin should be able to leave workspace'
+  );
+  t.true(
+    await leaveWorkspace(app, write.token.token, ws.id),
+    'write should be able to leave workspace'
+  );
+  t.true(
+    await leaveWorkspace(app, read.token.token, ws.id),
+    'read should be able to leave workspace'
+  );
+});
+
+test('should be able to invite by link', async t => {
+  const { app, permissions, quotaManager } = t.context;
+  const { createInviteLink, ws } = await init(app, 4);
+  const invite = await createInviteLink();
+  {
+    // invite link
+    const members: UserAuthedType[] = [];
+    await t.notThrowsAsync(async () => {
+      members.push(await invite('member3@affine.pro'));
+      members.push(await invite('member4@affine.pro'));
+    }, 'should not throw error even exceed member limit');
+    const [m3, m4] = members;
+
+    t.is(
+      await permissions.getWorkspaceMemberStatus(ws.id, m3.id),
+      WorkspaceMemberStatus.NeedMoreSeatAndReview,
+      'should not change status'
+    );
+    t.is(
+      await permissions.getWorkspaceMemberStatus(ws.id, m4.id),
+      WorkspaceMemberStatus.NeedMoreSeatAndReview,
+      'should not change status'
+    );
+
+    await quotaManager.updateWorkspaceConfig(ws.id, QuotaType.TeamPlanV1, {
+      memberLimit: 5,
+    });
+    await permissions.refreshSeatStatus(ws.id, 5);
+    t.is(
+      await permissions.getWorkspaceMemberStatus(ws.id, m3.id),
+      WorkspaceMemberStatus.UnderReview,
+      'should not change status'
+    );
+    t.is(
+      await permissions.getWorkspaceMemberStatus(ws.id, m4.id),
+      WorkspaceMemberStatus.NeedMoreSeatAndReview,
+      'should not change status'
+    );
+
+    await quotaManager.updateWorkspaceConfig(ws.id, QuotaType.TeamPlanV1, {
+      memberLimit: 6,
+    });
+    await permissions.refreshSeatStatus(ws.id, 6);
+    t.is(
+      await permissions.getWorkspaceMemberStatus(ws.id, m4.id),
+      WorkspaceMemberStatus.UnderReview,
+      'should not change status'
+    );
+  }
+});

packages/backend/server/tests/utils/feature.ts

@@ -0,0 +1,27 @@
+import { Injectable } from '@nestjs/common';
+import { PrismaClient, WorkspaceMemberStatus } from '@prisma/client';
+
+import { Permission } from '../../src/core/permission';
+import { UserType } from '../../src/core/user/types';
+
+@Injectable()
+export class WorkspaceResolverMock {
+  constructor(private readonly prisma: PrismaClient) {}
+
+  async createWorkspace(user: UserType, _init: null) {
+    const workspace = await this.prisma.workspace.create({
+      data: {
+        public: false,
+        permissions: {
+          create: {
+            type: Permission.Owner,
+            userId: user.id,
+            accepted: true,
+            status: WorkspaceMemberStatus.Accepted,
+          },
+        },
+      },
+    });
+    return workspace;
+  }
+}

packages/backend/server/tests/utils/invite.ts

@@ -3,13 +3,14 @@ import request from 'supertest';
 
 import type { InvitationType } from '../../src/core/workspaces';
 import { gql } from './common';
+import { PermissionEnum } from './utils';
 
 export async function inviteUser(
   app: INestApplication,
   token: string,
   workspaceId: string,
   email: string,
-  permission: string,
+  permission: PermissionEnum,
   sendInviteMail = false
 ): Promise<string> {
   const res = await request(app.getHttpServer())
@@ -24,18 +25,81 @@ export async function inviteUser(
           `,
     })
     .expect(200);
+  if (res.body.errors) {
+    throw new Error(res.body.errors[0].message);
+  }
   return res.body.data.invite;
 }
 
+export async function inviteUsers(
+  app: INestApplication,
+  token: string,
+  workspaceId: string,
+  emails: string[],
+  sendInviteMail = false
+): Promise<Array<{ email: string; inviteId?: string; sentSuccess?: boolean }>> {
+  const res = await request(app.getHttpServer())
+    .post(gql)
+    .auth(token, { type: 'bearer' })
+    .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
+    .send({
+      query: `
+          mutation inviteBatch($workspaceId: String!, $emails: [String!]!, $sendInviteMail: Boolean) {
+            inviteBatch(
+              workspaceId: $workspaceId
+              emails: $emails
+              sendInviteMail: $sendInviteMail
+            ) {
+              email
+              inviteId
+              sentSuccess
+            }
+          }
+          `,
+      variables: { workspaceId, emails, sendInviteMail },
+    })
+    .expect(200);
+  if (res.body.errors) {
+    throw new Error(res.body.errors[0].message);
+  }
+  return res.body.data.inviteBatch;
+}
+
+export async function inviteLink(
+  app: INestApplication,
+  token: string,
+  workspaceId: string,
+  expireTime: 'OneDay' | 'ThreeDays' | 'OneWeek' | 'OneMonth'
+): Promise<string> {
+  const res = await request(app.getHttpServer())
+    .post(gql)
+    .auth(token, { type: 'bearer' })
+    .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
+    .send({
+      query: `
+            mutation {
+              inviteLink(workspaceId: "${workspaceId}", expireTime: ${expireTime})
+            }
+          `,
+    })
+    .expect(200);
+  if (res.body.errors) {
+    throw new Error(res.body.errors[0].message);
+  }
+  return res.body.data.inviteLink;
+}
+
 export async function acceptInviteById(
   app: INestApplication,
   workspaceId: string,
   inviteId: string,
-  sendAcceptMail = false
+  sendAcceptMail = false,
+  token: string = ''
 ): Promise<boolean> {
   const res = await request(app.getHttpServer())
     .post(gql)
     .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
+    .auth(token, { type: 'bearer' })
     .send({
       query: `
             mutation {
@@ -44,6 +108,9 @@ export async function acceptInviteById(
           `,
     })
     .expect(200);
+  if (res.body.errors) {
+    throw new Error(res.body.errors[0].message);
+  }
   return res.body.data.acceptInviteById;
 }
 
@@ -65,6 +132,9 @@ export async function leaveWorkspace(
           `,
     })
     .expect(200);
+  if (res.body.errors) {
+    throw new Error(res.body.errors[0].message);
+  }
   return res.body.data.leaveWorkspace;
 }
 

packages/backend/server/tests/utils/user.ts

@@ -10,6 +10,8 @@ import { sessionUser } from '../../src/core/auth/service';
 import { UserService, type UserType } from '../../src/core/user';
 import { gql } from './common';
 
+export type UserAuthedType = UserType & { token: ClientTokenType };
+
 export async function internalSignIn(app: INestApplication, userId: string) {
   const auth = app.get(AuthService);
 
@@ -49,7 +51,7 @@ export async function signUp(
   email: string,
   password: string,
   autoVerifyEmail = true
-): Promise<UserType & { token: ClientTokenType }> {
+): Promise<UserAuthedType> {
   const user = await app.get(UserService).createUser({
     name,
     email,
@@ -176,7 +178,7 @@ export async function changeEmail(
   userToken: string,
   token: string,
   email: string
-): Promise<UserType & { token: ClientTokenType }> {
+): Promise<UserAuthedType> {
   const res = await request(app.getHttpServer())
     .post(gql)
     .auth(userToken, { type: 'bearer' })

packages/backend/server/tests/utils/utils.ts

@@ -14,6 +14,8 @@ import { UserFeaturesInit1698652531198 } from '../../src/data/migrations/1698652
 import { Config, GlobalExceptionFilter } from '../../src/fundamentals';
 import { GqlModule } from '../../src/fundamentals/graphql';
 
+export type PermissionEnum = 'Owner' | 'Admin' | 'Write' | 'Read';
+
 async function flushDB(client: PrismaClient) {
   const result: { tablename: string }[] =
     await client.$queryRaw`SELECT tablename

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

@@ -3,6 +3,7 @@ import request from 'supertest';
 
 import type { WorkspaceType } from '../../src/core/workspaces';
 import { gql } from './common';
+import { PermissionEnum } from './utils';
 
 export async function createWorkspace(
   app: INestApplication,
@@ -150,3 +151,32 @@ export async function revokePublicPage(
     .expect(200);
   return res.body.errors?.[0]?.message || res.body.data?.revokePublicPage;
 }
+
+export async function grantMember(
+  app: INestApplication,
+  token: string,
+  workspaceId: string,
+  userId: string,
+  permission: PermissionEnum
+) {
+  const res = await request(app.getHttpServer())
+    .post(gql)
+    .auth(token, { type: 'bearer' })
+    .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
+    .send({
+      query: `
+          mutation {
+            grantMember(
+              workspaceId: "${workspaceId}"
+              userId: "${userId}"
+              permission: ${permission}
+            )
+          }
+          `,
+    })
+    .expect(200);
+  if (res.body.errors) {
+    throw new Error(res.body.errors[0].message);
+  }
+  return res.body.data?.grantMember;
+}

packages/backend/server/tests/workspace/controller.spec.ts

@@ -1,7 +1,7 @@
 import { Readable } from 'node:stream';
 
 import { HttpStatus, INestApplication } from '@nestjs/common';
-import { PrismaClient } from '@prisma/client';
+import { PrismaClient, WorkspaceMemberStatus } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 import Sinon from 'sinon';
 import request from 'supertest';
@@ -182,6 +182,7 @@ test('should be able to get permission granted workspace', async t => {
       userId: u1.id,
       type: 1,
       accepted: true,
+      status: WorkspaceMemberStatus.Accepted,
     },
   });