[PB-6473]: fix/set quota during account provisioning

src/modules/account/account.service.spec.ts

@@ -3,6 +3,7 @@ import { Test, type TestingModule } from '@nestjs/testing';
 import { createMock, type DeepMocked } from '@golevelup/ts-vitest';
 import {
   ConflictException,
+  ForbiddenException,
   NotFoundException,
   UnprocessableEntityException,
 } from '@nestjs/common';
@@ -17,6 +18,8 @@ import { AddressRepository } from './repositories/address.repository.js';
 import { DomainRepository } from './repositories/domain.repository.js';
 import { MailAddressKeysRepository } from './repositories/mail-address-keys.repository.js';
 import { BridgeClient } from '../infrastructure/bridge/bridge.service.js';
+import { PaymentsService } from '../infrastructure/payments/payments.service.js';
+import type { Tier } from '../infrastructure/payments/payments.types.js';
 import {
   newMailAccountAttributes,
   newMailAddressKeyBundle,
@@ -35,6 +38,7 @@ describe('AccountService', () => {
   let domains: DeepMocked<DomainRepository>;
   let keys: DeepMocked<MailAddressKeysRepository>;
   let bridge: DeepMocked<BridgeClient>;
+  let payments: DeepMocked<PaymentsService>;
   let config: DeepMocked<ConfigService>;
 
   beforeEach(async () => {
@@ -51,6 +55,7 @@ describe('AccountService', () => {
     domains = module.get(DomainRepository);
     keys = module.get(MailAddressKeysRepository);
     bridge = module.get(BridgeClient);
+    payments = module.get(PaymentsService);
     config = module.get(ConfigService);
   });
 
@@ -411,6 +416,17 @@ describe('AccountService', () => {
       displayName: 'Alice Smith',
       keys: newMailAddressKeyBundle(),
     };
+    const PLAN_BYTES = 1_000_000_000;
+    const validTier: Tier = {
+      id: 't1',
+      label: 'standard',
+      productId: 'p1',
+      billingType: 'monthly',
+      featuresPerService: {
+        mail: { enabled: true, addressesPerUser: 3 },
+        drive: { enabled: true, maxSpaceBytes: PLAN_BYTES },
+      },
+    };
 
     it('when all inputs are valid, then creates account, address, provider link, and stalwart principal', async () => {
       const createdAccount = MailAccount.build(
@@ -437,6 +453,7 @@ describe('AccountService', () => {
       );
 
       const bucket = { id: 'bucket-1', name: createdAddressId };
+      payments.getUserTier.mockResolvedValue(validTier);
       domains.findByDomain.mockResolvedValue(domain);
       addresses.findByAddress.mockResolvedValue(null);
       accounts.findByUserId
@@ -476,6 +493,7 @@ describe('AccountService', () => {
           accountId: createdAccount.id,
           primaryAddress: params.address,
           displayName: params.displayName,
+          quota: PLAN_BYTES,
         }),
       );
       expect(keys.create).toHaveBeenCalledWith({
@@ -484,6 +502,42 @@ describe('AccountService', () => {
       });
     });
 
+    it('when the plan does not include mail, then throws ForbiddenException and creates nothing', async () => {
+      payments.getUserTier.mockResolvedValue({
+        ...validTier,
+        featuresPerService: {
+          drive: { enabled: true, maxSpaceBytes: PLAN_BYTES },
+        },
+      });
+      domains.findByDomain.mockResolvedValue(domain);
+      addresses.findByAddress.mockResolvedValue(null);
+      accounts.findByUserId.mockResolvedValue(null);
+
+      await expect(service.provisionAccount(params)).rejects.toThrow(
+        ForbiddenException,
+      );
+      expect(accounts.create).not.toHaveBeenCalled();
+      expect(provider.createAccount).not.toHaveBeenCalled();
+    });
+
+    it('when the plan has no drive storage allowance, then throws and never provisions an unlimited principal', async () => {
+      payments.getUserTier.mockResolvedValue({
+        ...validTier,
+        featuresPerService: {
+          mail: { enabled: true, addressesPerUser: 3 },
+        },
+      });
+      domains.findByDomain.mockResolvedValue(domain);
+      addresses.findByAddress.mockResolvedValue(null);
+      accounts.findByUserId.mockResolvedValue(null);
+
+      await expect(service.provisionAccount(params)).rejects.toThrow(
+        UnprocessableEntityException,
+      );
+      expect(accounts.create).not.toHaveBeenCalled();
+      expect(provider.createAccount).not.toHaveBeenCalled();
+    });
+
     it('when domain does not exist, then throws NotFoundException', async () => {
       domains.findByDomain.mockResolvedValue(null);
       addresses.findByAddress.mockResolvedValue(null);

src/modules/account/account.service.ts

@@ -1,6 +1,7 @@
 import { randomBytes } from 'node:crypto';
 import {
   ConflictException,
+  ForbiddenException,
   Injectable,
   Logger,
   NotFoundException,
@@ -9,6 +10,7 @@ import {
 import { ConfigService } from '@nestjs/config';
 import dayjs from 'dayjs';
 import { BridgeClient } from '../infrastructure/bridge/bridge.service.js';
+import { PaymentsService } from '../infrastructure/payments/payments.service.js';
 import { MailNotSetupException } from '../provisioning/mail-not-setup.exception.js';
 import { AccountProvider } from './account-provider.port.js';
 import { MailAccount, MailAccountState } from './domain/mail-account.domain.js';
@@ -49,6 +51,7 @@ export class AccountService {
     private readonly domains: DomainRepository,
     private readonly keys: MailAddressKeysRepository,
     private readonly bridge: BridgeClient,
+    private readonly payments: PaymentsService,
     private readonly config: ConfigService,
   ) {}
 
@@ -184,12 +187,19 @@ export class AccountService {
     displayName: string;
     keys: MailAddressKeyBundle;
   }): Promise<MailAccount> {
-    const [domainRecord, existingAddress, existingAccount] = await Promise.all([
-      this.domains.findByDomain(params.domain),
-      this.addresses.findByAddress(params.address),
-      this.accounts.findByUserId(params.userId),
-    ]);
-
+    const [tier, domainRecord, existingAddress, existingAccount] =
+      await Promise.all([
+        this.payments.getUserTier(params.userId),
+        this.domains.findByDomain(params.domain),
+        this.addresses.findByAddress(params.address),
+        this.accounts.findByUserId(params.userId),
+      ]);
+
+    if (!tier.featuresPerService.mail?.enabled) {
+      throw new ForbiddenException(
+        'Mail access is not available for your current plan',
+      );
+    }
     if (!domainRecord) {
       throw new NotFoundException(`Domain '${params.domain}' not found`);
     }
@@ -202,6 +212,13 @@ export class AccountService {
       );
     }
 
+    const quota = tier.featuresPerService.drive?.maxSpaceBytes;
+    if (!quota || quota <= 0) {
+      throw new UnprocessableEntityException(
+        `Cannot provision mail account for '${params.userId}': plan has no drive storage allowance`,
+      );
+    }
+
     let account: MailAccount;
     try {
       account = await this.accounts.create({
@@ -246,6 +263,7 @@ export class AccountService {
         primaryAddress: params.address,
         displayName: params.displayName,
         password,
+        quota,
       });
     } catch (error) {
       await this.accounts.delete(account.id);

src/modules/account/user.controller.spec.ts

@@ -1,33 +1,19 @@
 import { describe, it, expect, beforeEach } from 'vitest';
 import { Test, type TestingModule } from '@nestjs/testing';
 import { createMock, type DeepMocked } from '@golevelup/ts-vitest';
-import { ForbiddenException } from '@nestjs/common';
 import { UserController } from './user.controller.js';
 import { AccountService, type MailAccountStatus } from './account.service.js';
 import { MailAccountState, MailAccount } from './domain/mail-account.domain.js';
-import { PaymentsService } from '../infrastructure/payments/payments.service.js';
 import {
   newMailAccountAttributes,
   newMailAddressKeyBundle,
   newUserPayload,
 } from '../../../test/fixtures.js';
 import type { CreateMailAccountDto } from './dto/create-mail-account.dto.js';
-import type { Tier } from '../infrastructure/payments/payments.types.js';
-
-const tierWith = (mailEnabled: boolean): Tier => ({
-  id: 't1',
-  label: 'ultimate',
-  productId: 'p1',
-  billingType: 'monthly',
-  featuresPerService: {
-    mail: { enabled: mailEnabled, addressesPerUser: 3 },
-  },
-});
 
 describe('UserController', () => {
   let controller: UserController;
   let accountService: DeepMocked<AccountService>;
-  let payments: DeepMocked<PaymentsService>;
 
   const buildDto = (): CreateMailAccountDto => ({
     address: 'alice',
@@ -45,7 +31,6 @@ describe('UserController', () => {
 
     controller = module.get(UserController);
     accountService = module.get(AccountService);
-    payments = module.get(PaymentsService);
   });
 
   describe('getMailAccount', () => {
@@ -127,23 +112,12 @@ describe('UserController', () => {
   });
 
   describe('createMailAccount', () => {
-    it('when tier disables mail, then throws ForbiddenException', async () => {
-      const user = newUserPayload();
-      payments.getUserTier.mockResolvedValue(tierWith(false));
-
-      await expect(
-        controller.createMailAccount(user, buildDto()),
-      ).rejects.toThrow(ForbiddenException);
-      expect(accountService.provisionAccount).not.toHaveBeenCalled();
-    });
-
-    it('when all checks pass, then provisions and returns address', async () => {
+    it('when provisioning succeeds, then returns the address', async () => {
       const user = newUserPayload();
       const dto = buildDto();
       const account = MailAccount.build(
         newMailAccountAttributes({ userId: user.uuid }),
       );
-      payments.getUserTier.mockResolvedValue(tierWith(true));
       accountService.provisionAccount.mockResolvedValue(account);
 
       const result = await controller.createMailAccount(user, dto);

src/modules/account/user.controller.ts

@@ -1,7 +1,6 @@
 import {
   Body,
   Controller,
-  ForbiddenException,
   Get,
   HttpCode,
   HttpStatus,
@@ -23,7 +22,6 @@ import {
 import { GetMailAccountKeysDto } from './dto/get-mail-account-keys.dto.js';
 import { User } from '../auth/decorators/user.decorator.js';
 import type { UserPayload } from '../auth/jwt-payload.dto.js';
-import { PaymentsService } from '../infrastructure/payments/payments.service.js';
 import { AccountService } from './account.service.js';
 import { CreateMailAccountDto } from './dto/create-mail-account.dto.js';
 import { MailAccountGuard } from '../provisioning/provisioning.guard.js';
@@ -40,10 +38,7 @@ import {
 export class UserController {
   private readonly logger = new Logger(UserController.name);
 
-  constructor(
-    private readonly accountService: AccountService,
-    private readonly payments: PaymentsService,
-  ) {}
+  constructor(private readonly accountService: AccountService) {}
 
   @Get('me/mail-account')
   @UseGuards(MailAccountGuard)
@@ -75,13 +70,6 @@ export class UserController {
     @User() user: UserPayload,
     @Body() dto: CreateMailAccountDto,
   ): Promise<CreateMailAccountResponseDto> {
-    const tier = await this.payments.getUserTier(user.uuid);
-    if (!tier.featuresPerService.mail?.enabled) {
-      throw new ForbiddenException(
-        'Mail access is not available for your current plan',
-      );
-    }
-
     const fullAddress = `${dto.address}@${dto.domain}`;
 
     const account = await this.accountService.provisionAccount({