chore: add quota usage e2e tests (#2438)

* chore: add quota usage e2e tests

* chore: add moving from shared to shared vault quota e2e test

* chore: fix test suite name

* chore: fix awaiting for notifications to propagate

* chore: fix awaiting for notifications processing

packages/snjs/mocha/TestRegistry/VaultTests.js

@@ -22,5 +22,6 @@ export const VaultTests = {
     'vaults/key-rotation.test.js',
     'vaults/files.test.js',
     'vaults/limits.test.js',
+    'vaults/quota.test.js',
   ],
 }

packages/snjs/mocha/lib/AppContext.js

@@ -149,6 +149,10 @@ export class AppContext {
     return this.application.asymmetric
   }
 
+  get notifications() {
+    return this.application.dependencies.get(TYPES.NotificationService)
+  }
+
   get keyPair() {
     return this.application.dependencies.get(TYPES.GetKeyPairs).execute().getValue().encryption
   }
@@ -476,11 +480,6 @@ export class AppContext {
     })
   }
 
-  resolveWhenUserMessagesProcessingCompletes() {
-    const objectToSpy = this.application.dependencies.get(TYPES.NotificationService)
-    return this.resolveWhenAsyncFunctionCompletes(objectToSpy, 'handleReceivedNotifications')
-  }
-
   resolveWhenAllInboundAsymmetricMessagesAreDeleted() {
     const objectToSpy = this.application.dependencies.get(TYPES.AsymmetricMessageServer)
     return this.resolveWhenAsyncFunctionCompletes(objectToSpy, 'deleteAllInboundMessages')
@@ -658,8 +657,8 @@ export class AppContext {
     return this.application.sessions.user.uuid
   }
 
-  sleep(seconds) {
-    return Utils.sleep(seconds)
+  sleep(seconds, reason = undefined) {
+    return Utils.sleep(seconds, reason)
   }
 
   anticipateConsoleError(message, _reason) {
@@ -670,12 +669,25 @@ export class AppContext {
     return Utils.awaitPromiseOrThrow(promise, maxWait, reason)
   }
 
+  awaitPromiseOrDoNothing(promise, maxWait = 2.0, reason = 'Awaiting promise timed out; No description provided') {
+    return Utils.awaitPromiseOrDoNothing(promise, maxWait, reason)
+  }
+
   async activatePaidSubscriptionForUser(options = {}) {
     const dateInAnHour = new Date()
     dateInAnHour.setHours(dateInAnHour.getHours() + 1)
 
     options.expiresAt = options.expiresAt || dateInAnHour
     options.subscriptionPlanName = options.subscriptionPlanName || 'PRO_PLAN'
+    let uploadBytesLimit = -1
+    switch (options.subscriptionPlanName) {
+      case 'PLUS_PLAN':
+        uploadBytesLimit = 104_857_600
+        break
+      case 'PRO_PLAN':
+        uploadBytesLimit = 107_374_182_400
+        break
+    }
 
     try {
       await Events.publishMockedEvent('SUBSCRIPTION_PURCHASED', {
@@ -694,16 +706,16 @@ export class AppContext {
         payAmount: 59.0,
       })
 
-      await Utils.sleep(2, 'Waiting for premium features to be activated')
+      await this.sleep(2, 'Waiting for premium features to be activated')
     } catch (error) {
       console.warn(
         `Mock events service not available. You are probably running a test suite for home server: ${error.message}`,
       )
 
       try {
-        await HomeServer.activatePremiumFeatures(this.email, options.subscriptionPlanName, options.expiresAt)
+        await HomeServer.activatePremiumFeatures(this.email, options.subscriptionPlanName, options.expiresAt, uploadBytesLimit)
 
-        await Utils.sleep(1, 'Waiting for premium features to be activated')
+        await this.sleep(1, 'Waiting for premium features to be activated')
       } catch (error) {
         console.warn(
           `Home server not available. You are probably running a test suite for self hosted setup: ${error.message}`,

packages/snjs/mocha/lib/Files.js

@@ -1,5 +1,9 @@
-export async function uploadFile(fileService, buffer, name, ext, chunkSize, vault) {
-  const operation = await fileService.beginNewFileUpload(buffer.byteLength, vault)
+export async function uploadFile(fileService, buffer, name, ext, chunkSize, vault, options = {}) {
+  const byteLength = options.byteLengthOverwrite || buffer.byteLength
+  const operation = await fileService.beginNewFileUpload(byteLength, vault)
+  if (isClientDisplayableError(operation)) {
+    return operation
+  }
 
   let chunkId = 1
   for (let i = 0; i < buffer.length; i += chunkSize) {

packages/snjs/mocha/lib/HomeServer.js

@@ -1,6 +1,6 @@
 import * as Defaults from './Defaults.js'
 
-export async function activatePremiumFeatures(username, subscriptionPlanName, endsAt) {
+export async function activatePremiumFeatures(username, subscriptionPlanName, endsAt, uploadBytesLimit) {
   await fetch(`${Defaults.getDefaultHost()}/e2e/activate-premium`, {
     method: 'POST',
     headers: {
@@ -11,6 +11,7 @@ export async function activatePremiumFeatures(username, subscriptionPlanName, en
       username,
       subscriptionPlanName,
       endsAt,
+      uploadBytesLimit,
     }),
   })
 }

packages/snjs/mocha/lib/Utils.js

@@ -53,3 +53,23 @@ export async function awaitPromiseOrThrow(promise, maxWait, reason) {
     return result
   })
 }
+
+export async function awaitPromiseOrDoNothing(promise, maxWait, reason) {
+  let timer = undefined
+
+  // Create a promise that resolves in <maxWait> milliseconds
+  const timeout = new Promise((resolve, reject) => {
+    timer = setTimeout(() => {
+      clearTimeout(timer)
+      const message = reason || `Promise timed out after ${maxWait} milliseconds: ${reason}`
+      console.warn(message)
+      resolve()
+    }, maxWait * 1000)
+  })
+
+  // Returns a race between our timeout and the passed in promise
+  return Promise.race([promise, timeout]).then((result) => {
+    clearTimeout(timer)
+    return result
+  })
+}

packages/snjs/mocha/lib/VaultsContext.js

@@ -34,6 +34,24 @@ export class VaultsContext extends AppContext {
     await this.awaitPromiseOrThrow(promise, undefined, 'Waiting for keypair change message to process')
   }
 
+  async syncAndAwaitNotificationsProcessing() {
+    await this.sleep(0.25, 'Waiting for notifications to propagate')
+
+    const promise = this.resolveWhenAsyncFunctionCompletes(this.notifications, 'handleReceivedNotifications')
+
+    await this.sync()
+
+    await this.awaitPromiseOrDoNothing(
+      promise,
+      0.25,
+      'Waiting for notifications timed out. Notifications might have been processed in previous sync.'
+    )
+
+    if (this.notifications['handleReceivedNotifications'].restore) {
+      this.notifications['handleReceivedNotifications'].restore()
+    }
+  }
+
   async syncAndAwaitMessageProcessing() {
     const promise = this.resolveWhenAsyncFunctionCompletes(this.asymmetric, 'handleRemoteReceivedAsymmetricMessages')
 

packages/snjs/mocha/vaults/quota.test.js

@@ -0,0 +1,213 @@
+import * as Factory from '../lib/factory.js'
+import * as Files from '../lib/Files.js'
+import * as Collaboration from '../lib/Collaboration.js'
+
+chai.use(chaiAsPromised)
+const expect = chai.expect
+
+describe('shared vault quota', function () {
+  this.timeout(Factory.TwentySecondTimeout)
+
+  let context
+
+  beforeEach(async function () {
+    localStorage.clear()
+
+    context = await Factory.createVaultsContextWithRealCrypto()
+
+    await context.launch()
+    await context.register()
+  })
+
+  afterEach(async function () {
+    await context.deinit()
+    localStorage.clear()
+    sinon.restore()
+    context = undefined
+  })
+
+  describe('using own quota', function () {
+    it('should utilize my own quota when I am uploading to my shared vault', async () => {
+      await context.activatePaidSubscriptionForUser()
+
+      const sharedVault = await Collaboration.createSharedVault(context)
+      const response = await fetch('/mocha/assets/small_file.md')
+      const buffer = new Uint8Array(await response.arrayBuffer())
+      await Files.uploadFile(context.files, buffer, 'my-file', 'md', 1000, sharedVault)
+
+      await context.syncAndAwaitNotificationsProcessing()
+
+      const updatedVault = context.vaults.getVault({ keySystemIdentifier: sharedVault.systemIdentifier })
+      expect(updatedVault.sharing.fileBytesUsed).to.equal(1374)
+
+      const bytesUsedSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+bytesUsedSetting).to.equal(1374)
+    })
+
+    it('should not allow me to upload a file that exceeds my quota', async () => {
+      await context.activatePaidSubscriptionForUser()
+
+      const bytesLimitSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesLimit).getValue(),
+      )
+      expect(+bytesLimitSetting).to.equal(107_374_182_400)
+
+      const sharedVault = await Collaboration.createSharedVault(context)
+      const response = await fetch('/mocha/assets/small_file.md')
+      const buffer = new Uint8Array(await response.arrayBuffer())
+      const result = await Files.uploadFile(context.files, buffer, 'my-file', 'md', 1000, sharedVault, { byteLengthOverwrite: 107_374_182_401 })
+
+      expect(isClientDisplayableError(result)).to.be.true
+
+      const bytesUsedSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+bytesUsedSetting).to.equal(0)
+    })
+
+    it('should utilize my own quota when I am moving a user file to my vault', async () => {
+      await context.activatePaidSubscriptionForUser()
+
+      const response = await fetch('/mocha/assets/small_file.md')
+      const buffer = new Uint8Array(await response.arrayBuffer())
+
+      const uploadedFile = await Files.uploadFile(context.files, buffer, 'my-file', 'md', 1000)
+
+      const sharedVault = await Collaboration.createSharedVault(context)
+      await context.vaults.moveItemToVault(sharedVault, uploadedFile)
+
+      await context.syncAndAwaitNotificationsProcessing()
+
+      const updatedVault = context.vaults.getVault({ keySystemIdentifier: sharedVault.systemIdentifier })
+      expect(updatedVault.sharing.fileBytesUsed).to.equal(1374)
+
+      const bytesUsedSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+bytesUsedSetting).to.equal(1374)
+    })
+  })
+
+  describe('using contact quota', function () {
+    it('should utilize my quota when contact is uploading to my shared vault', async () => {
+      await context.activatePaidSubscriptionForUser()
+
+      const { sharedVault, contactContext, deinitContactContext } =
+        await Collaboration.createSharedVaultWithAcceptedInvite(context)
+      await contactContext.activatePaidSubscriptionForUser()
+
+      const response = await fetch('/mocha/assets/small_file.md')
+      const buffer = new Uint8Array(await response.arrayBuffer())
+
+      await Files.uploadFile(contactContext.files, buffer, 'my-file', 'md', 1000, sharedVault)
+
+      await context.syncAndAwaitNotificationsProcessing()
+      await contactContext.syncAndAwaitNotificationsProcessing()
+
+      const updatedVault = context.vaults.getVault({ keySystemIdentifier: sharedVault.systemIdentifier })
+      expect(updatedVault.sharing.fileBytesUsed).to.equal(1374)
+
+      const myBytesUsedSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+myBytesUsedSetting).to.equal(1374)
+
+      const contactBytesUsedSetting = await contactContext.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+contactBytesUsedSetting).to.equal(0)
+
+      await deinitContactContext()
+    })
+
+    it('should not allow my contact to upload a file that exceeds my quota', async () => {
+      await context.activatePaidSubscriptionForUser({ subscriptionPlanName: 'PLUS_PLAN' })
+
+      const myBytesLimitSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesLimit).getValue(),
+      )
+      expect(+myBytesLimitSetting).to.equal(104_857_600)
+
+      const { sharedVault, contactContext, deinitContactContext } =
+        await Collaboration.createSharedVaultWithAcceptedInvite(context)
+      await contactContext.activatePaidSubscriptionForUser({ subscriptionPlanName: 'PRO_PLAN' })
+
+      const contactBytesLimitSetting = await contactContext.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesLimit).getValue(),
+      )
+      expect(+contactBytesLimitSetting).to.equal(107_374_182_400)
+
+      const response = await fetch('/mocha/assets/small_file.md')
+      const buffer = new Uint8Array(await response.arrayBuffer())
+      const result = await Files.uploadFile(context.files, buffer, 'my-file', 'md', 1000, sharedVault, { byteLengthOverwrite: 104_857_601 })
+
+      expect(isClientDisplayableError(result)).to.be.true
+
+      const bytesUsedSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+bytesUsedSetting).to.equal(0)
+
+      await deinitContactContext()
+    })
+
+    it('should utilize my quota when my contact is moving a shared file from contact vault to my vault', async () => {
+      await context.activatePaidSubscriptionForUser()
+
+      const { sharedVault, contactContext, deinitContactContext } =
+        await Collaboration.createSharedVaultWithAcceptedInvite(context)
+      await contactContext.activatePaidSubscriptionForUser()
+
+      const secondVault = await Collaboration.createSharedVault(contactContext)
+
+      const response = await fetch('/mocha/assets/small_file.md')
+      const buffer = new Uint8Array(await response.arrayBuffer())
+
+      const uploadedFile = await Files.uploadFile(contactContext.files, buffer, 'my-file', 'md', 1000, secondVault)
+
+      await context.syncAndAwaitNotificationsProcessing()
+      await contactContext.syncAndAwaitNotificationsProcessing()
+
+      let updatedSharedVault = context.vaults.getVault({ keySystemIdentifier: sharedVault.systemIdentifier })
+      expect(updatedSharedVault.sharing.fileBytesUsed).to.equal(0)
+
+      let updatedSecondVault = contactContext.vaults.getVault({ keySystemIdentifier: secondVault.systemIdentifier })
+      expect(updatedSecondVault.sharing.fileBytesUsed).to.equal(1374)
+
+      let myBytesUsedSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+myBytesUsedSetting).to.equal(0)
+
+      let contactBytesUsedSetting = await contactContext.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+contactBytesUsedSetting).to.equal(1374)
+
+      await contactContext.vaults.moveItemToVault(sharedVault, uploadedFile)
+
+      await context.syncAndAwaitNotificationsProcessing()
+      await contactContext.syncAndAwaitNotificationsProcessing()
+
+      updatedSharedVault = context.vaults.getVault({ keySystemIdentifier: sharedVault.systemIdentifier })
+      expect(updatedSharedVault.sharing.fileBytesUsed).to.equal(1374)
+
+      updatedSecondVault = contactContext.vaults.getVault({ keySystemIdentifier: secondVault.systemIdentifier })
+      expect(updatedSecondVault.sharing.fileBytesUsed).to.equal(0)
+
+      myBytesUsedSetting = await context.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+myBytesUsedSetting).to.equal(1374)
+
+      contactBytesUsedSetting = await contactContext.application.settings.getSubscriptionSetting(
+        SettingName.create(SettingName.NAMES.FileUploadBytesUsed).getValue(),
+      )
+      expect(+contactBytesUsedSetting).to.equal(0)
+
+      await deinitContactContext()
+    })
+  })
+})

packages/snjs/mocha/vaults/shared_vaults.test.js

@@ -57,9 +57,7 @@ describe('shared vaults', function () {
     const result = await context.vaultUsers.removeUserFromSharedVault(sharedVault, contactContext.userUuid)
     expect(result.isFailed()).to.be.false
 
-    const promise = contactContext.resolveWhenUserMessagesProcessingCompletes()
-    await contactContext.sync()
-    await promise
+    await contactContext.syncAndAwaitNotificationsProcessing()
 
     expect(contactContext.vaults.getVault({ keySystemIdentifier: sharedVault.systemIdentifier })).to.be.undefined
     expect(contactContext.keys.getPrimaryKeySystemRootKey(sharedVault.systemIdentifier)).to.be.undefined
@@ -84,9 +82,7 @@ describe('shared vaults', function () {
 
     expect(result).to.be.undefined
 
-    const promise = contactContext.resolveWhenUserMessagesProcessingCompletes()
-    await contactContext.sync()
-    await promise
+    await contactContext.syncAndAwaitNotificationsProcessing()
 
     expect(contactContext.vaults.getVault({ keySystemIdentifier: sharedVault.systemIdentifier })).to.be.undefined
     expect(contactContext.keys.getPrimaryKeySystemRootKey(sharedVault.systemIdentifier)).to.be.undefined