From 9e7a2fcf0ec8b64c09064baba8386955983a6268 Mon Sep 17 00:00:00 2001
From: EYHN <cneyhn@gmail.com>
Date: Fri, 12 Apr 2024 06:45:18 +0000
Subject: [PATCH] feat(server): add pro quota to dev user (#6532)

---
 packages/backend/server/src/core/auth/index.ts     |  3 ++-
 packages/backend/server/src/core/auth/service.ts   | 12 ++++++++++--
 packages/backend/server/tests/auth/service.spec.ts |  3 ++-
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/packages/backend/server/src/core/auth/index.ts b/packages/backend/server/src/core/auth/index.ts
index 318075f74..b840a5108 100644
--- a/packages/backend/server/src/core/auth/index.ts
+++ b/packages/backend/server/src/core/auth/index.ts
@@ -1,6 +1,7 @@
 import { Module } from '@nestjs/common';
 
 import { FeatureModule } from '../features';
+import { QuotaModule } from '../quota';
 import { UserModule } from '../user';
 import { AuthController } from './controller';
 import { AuthResolver } from './resolver';
@@ -8,7 +9,7 @@ import { AuthService } from './service';
 import { TokenService, TokenType } from './token';
 
 @Module({
-  imports: [FeatureModule, UserModule],
+  imports: [FeatureModule, UserModule, QuotaModule],
   providers: [AuthService, AuthResolver, TokenService],
   exports: [AuthService],
   controllers: [AuthController],
diff --git a/packages/backend/server/src/core/auth/service.ts b/packages/backend/server/src/core/auth/service.ts
index 86adf281b..54021c6bc 100644
--- a/packages/backend/server/src/core/auth/service.ts
+++ b/packages/backend/server/src/core/auth/service.ts
@@ -11,6 +11,8 @@ import { assign, omit } from 'lodash-es';
 
 import { Config, CryptoHelper, MailService } from '../../fundamentals';
 import { FeatureManagementService } from '../features/management';
+import { QuotaService } from '../quota/service';
+import { QuotaType } from '../quota/types';
 import { UserService } from '../user/service';
 import type { CurrentUser } from './current-user';
 
@@ -68,15 +70,21 @@ export class AuthService implements OnApplicationBootstrap {
     private readonly db: PrismaClient,
     private readonly mailer: MailService,
     private readonly feature: FeatureManagementService,
+    private readonly quota: QuotaService,
     private readonly user: UserService,
     private readonly crypto: CryptoHelper
   ) {}
 
   async onApplicationBootstrap() {
     if (this.config.node.dev) {
-      await this.signUp('Dev User', 'dev@affine.pro', 'dev').catch(() => {
+      try {
+        const devUser = await this.signUp('Dev User', 'dev@affine.pro', 'dev');
+        if (devUser) {
+          await this.quota.switchUserQuota(devUser?.id, QuotaType.ProPlanV1);
+        }
+      } catch (e) {
         // ignore
-      });
+      }
     }
   }
 
diff --git a/packages/backend/server/tests/auth/service.spec.ts b/packages/backend/server/tests/auth/service.spec.ts
index 55fbc24dd..3d4db7b37 100644
--- a/packages/backend/server/tests/auth/service.spec.ts
+++ b/packages/backend/server/tests/auth/service.spec.ts
@@ -5,6 +5,7 @@ import ava, { TestFn } from 'ava';
 import { CurrentUser } from '../../src/core/auth';
 import { AuthService, parseAuthUserSeqNum } from '../../src/core/auth/service';
 import { FeatureModule } from '../../src/core/features';
+import { QuotaModule } from '../../src/core/quota';
 import { UserModule, UserService } from '../../src/core/user';
 import { createTestingModule } from '../utils';
 
@@ -18,7 +19,7 @@ const test = ava as TestFn<{
 
 test.beforeEach(async t => {
   const m = await createTestingModule({
-    imports: [FeatureModule, UserModule],
+    imports: [QuotaModule, FeatureModule, UserModule],
     providers: [AuthService],
   });