Merge pull request #21472 from overleaf/em-hackathon-mongo-mocks-docker

Do not mock Mongo in unit tests

GitOrigin-RevId: 7a200a4ddc8f91b14e96cf02cb4873c51fc3489a

services/web/Makefile

@@ -67,18 +67,20 @@ test_module: test_unit_module test_acceptance_module
 #
 
 test_unit: test_unit_all
+test_unit_all: export COMPOSE_PROJECT_NAME=unit_test_all_$(BUILD_DIR_NAME)
 test_unit_all:
-	COMPOSE_PROJECT_NAME=unit_test_all_$(BUILD_DIR_NAME) $(DOCKER_COMPOSE) run --rm test_unit npm run test:unit:all
-	COMPOSE_PROJECT_NAME=unit_test_all_$(BUILD_DIR_NAME) $(DOCKER_COMPOSE) down -v -t 0
+	$(DOCKER_COMPOSE) run --rm test_unit npm run test:unit:all
+	$(DOCKER_COMPOSE) down -v -t 0
 
+test_unit_all_silent: export COMPOSE_PROJECT_NAME=unit_test_all_$(BUILD_DIR_NAME)
 test_unit_all_silent:
-	COMPOSE_PROJECT_NAME=unit_test_all_$(BUILD_DIR_NAME) $(DOCKER_COMPOSE) run --rm test_unit npm run test:unit:all:silent
-	COMPOSE_PROJECT_NAME=unit_test_all_$(BUILD_DIR_NAME) $(DOCKER_COMPOSE) down -v -t 0
+	$(DOCKER_COMPOSE) run --rm test_unit npm run test:unit:all:silent
+	$(DOCKER_COMPOSE) down -v -t 0
 
+test_unit_app: export COMPOSE_PROJECT_NAME=unit_test_$(BUILD_DIR_NAME)
 test_unit_app:
-	COMPOSE_PROJECT_NAME=unit_test_$(BUILD_DIR_NAME) $(DOCKER_COMPOSE) down -v -t 0
-	COMPOSE_PROJECT_NAME=unit_test_$(BUILD_DIR_NAME) $(DOCKER_COMPOSE) run --name unit_test_$(BUILD_DIR_NAME) --rm test_unit
-	COMPOSE_PROJECT_NAME=unit_test_$(BUILD_DIR_NAME) $(DOCKER_COMPOSE) down -v -t 0
+	$(DOCKER_COMPOSE) run --name unit_test_$(BUILD_DIR_NAME) --rm test_unit
+	$(DOCKER_COMPOSE) down -v -t 0
 
 TEST_SUITES = $(sort $(filter-out \
 	$(wildcard test/unit/src/helpers/*), \
@@ -101,7 +103,6 @@ test_unit_app_parallel_gnu_make: $(TEST_SUITES)
 test_unit_app_parallel_gnu_make_docker: export COMPOSE_PROJECT_NAME = \
 	unit_test_parallel_make_$(BUILD_DIR_NAME)
 test_unit_app_parallel_gnu_make_docker:
-	$(DOCKER_COMPOSE) down -v -t 0
 	$(DOCKER_COMPOSE) run --rm test_unit \
 		make test_unit_app_parallel_gnu_make --output-sync -j $(J)
 	$(DOCKER_COMPOSE) down -v -t 0
@@ -157,16 +158,6 @@ test_frontend_ct_editor:
 # Acceptance tests
 #
 
-# Keep in sync with TEST_ACCEPTANCE_MONGO_INIT in Makefile.module
-TEST_ACCEPTANCE_MONGO_INIT := \
-	$(DOCKER_COMPOSE) up -d mongo; \
-	$(DOCKER_COMPOSE) exec -T mongo sh -c ' \
-		while ! mongosh --eval "db.version()" > /dev/null; do \
-			echo "Waiting for Mongo..."; \
-			sleep 1; \
-		done; \
-		mongosh --eval "rs.initiate({ _id: \"overleaf\", members: [ { _id: 0, host: \"mongo:27017\" } ] })"'
-
 test_acceptance: test_acceptance_app test_acceptance_modules
 test_acceptance_saas: test_acceptance_app_saas test_acceptance_modules_merged_saas
 test_acceptance_server_ce: test_acceptance_app_server_ce test_acceptance_modules_merged_server_ce
@@ -186,8 +177,6 @@ test_acceptance_app_server_pro: export COMPOSE_PROJECT_NAME=acceptance_test_serv
 test_acceptance_app_server_pro: export OVERLEAF_CONFIG=$(CFG_SERVER_PRO)
 
 $(TEST_ACCEPTANCE_APP):
-	$(DOCKER_COMPOSE) down -v -t 0
-	$(TEST_ACCEPTANCE_MONGO_INIT)
 	$(DOCKER_COMPOSE) run --rm test_acceptance
 	$(DOCKER_COMPOSE) down -v -t 0
 
@@ -335,8 +324,6 @@ TEST_ACCEPTANCE_MODULES_MERGED_VARIANTS = \
 	test_acceptance_modules_merged_server_pro \
 
 $(TEST_ACCEPTANCE_MODULES_MERGED_VARIANTS):
-	$(DOCKER_COMPOSE) down -v -t 0
-	$(TEST_ACCEPTANCE_MONGO_INIT)
 	$(DOCKER_COMPOSE) run --rm test_acceptance make test_acceptance_modules_merged_inner
 	$(DOCKER_COMPOSE) down -v -t 0
 
@@ -358,8 +345,6 @@ test_acceptance_modules_merged_saas_4: export COMPOSE_PROJECT_NAME = \
 $(TEST_ACCEPTANCE_MODULES_MERGED_SPLIT_SAAS): export BASE_CONFIG = $(CFG_SAAS)
 
 $(TEST_ACCEPTANCE_MODULES_MERGED_SPLIT_SAAS): test_acceptance_modules_merged_saas_%:
-	$(DOCKER_COMPOSE) down -v -t 0
-	$(TEST_ACCEPTANCE_MONGO_INIT)
 	$(DOCKER_COMPOSE) run --rm test_acceptance make test_acceptance_modules_merged_inner_$*
 	$(DOCKER_COMPOSE) down -v -t 0
 

services/web/Makefile.module

@@ -19,31 +19,14 @@ DOCKER_COMPOSE := cd ../../ && \
 	MOCHA_GREP=${MOCHA_GREP} \
 	docker compose ${DOCKER_COMPOSE_FLAGS}
 
-DOCKER_COMPOSE_TEST_ACCEPTANCE := \
-	export COMPOSE_PROJECT_NAME=acceptance_test_$(BUILD_DIR_NAME)_$(MODULE_NAME) \
-	&& $(DOCKER_COMPOSE)
-
-DOCKER_COMPOSE_TEST_UNIT := \
-	export COMPOSE_PROJECT_NAME=unit_test_$(BUILD_DIR_NAME)_$(MODULE_NAME) \
-	&& $(DOCKER_COMPOSE)
-
-# Keep in sync with TEST_ACCEPTANCE_MONGO_INIT in Makefile
-TEST_ACCEPTANCE_MONGO_INIT := \
-	$(DOCKER_COMPOSE_TEST_ACCEPTANCE) up -d mongo; \
-	$(DOCKER_COMPOSE_TEST_ACCEPTANCE) exec -T mongo sh -c ' \
-		while ! mongosh --eval "db.version()" > /dev/null; do \
-			echo "Waiting for Mongo..."; \
-			sleep 1; \
-		done; \
-		mongosh --eval "rs.initiate({ _id: \"overleaf\", members: [ { _id: 0, host: \"mongo:27017\" } ] })"'
-
 ifeq (,$(wildcard test/unit))
 test_unit:
 
 else
+test_unit: export COMPOSE_PROJECT_NAME=unit_test_$(BUILD_DIR_NAME)_$(MODULE_NAME)
 test_unit:
-	${DOCKER_COMPOSE_TEST_UNIT} run --rm test_unit npm -q run test:unit:run_dir -- ${MOCHA_ARGS} $(MODULE_DIR)/test/unit/src
-	${DOCKER_COMPOSE_TEST_UNIT} down
+	${DOCKER_COMPOSE} run --rm test_unit npm -q run test:unit:run_dir -- ${MOCHA_ARGS} $(MODULE_DIR)/test/unit/src
+	${DOCKER_COMPOSE} down
 
 endif
 
@@ -66,17 +49,18 @@ test_acceptance_saas: export BASE_CONFIG = $(CFG_SAAS)
 test_acceptance_server_ce: export BASE_CONFIG = $(CFG_SERVER_CE)
 test_acceptance_server_pro: export BASE_CONFIG = $(CFG_SERVER_PRO)
 
+$(ALL_TEST_ACCEPTANCE_VARIANTS): export COMPOSE_PROJECT_NAME=acceptance_test_$(BUILD_DIR_NAME)_$(MODULE_NAME)
 $(ALL_TEST_ACCEPTANCE_VARIANTS):
 	$(MAKE) --no-print-directory clean_test_acceptance
-	$(TEST_ACCEPTANCE_MONGO_INIT)
-	${DOCKER_COMPOSE_TEST_ACCEPTANCE} run --rm test_acceptance npm -q run test:acceptance:run_dir -- ${MOCHA_ARGS} $(MODULE_DIR)/test/acceptance/src
+	${DOCKER_COMPOSE} run --rm test_acceptance npm -q run test:acceptance:run_dir -- ${MOCHA_ARGS} $(MODULE_DIR)/test/acceptance/src
 	$(MAKE) --no-print-directory clean_test_acceptance
 
 test_acceptance_merged_inner:
 	cd ../../ && \
 			npm -q run test:acceptance:run_dir -- ${MOCHA_ARGS} $(MODULE_DIR)/test/acceptance/src
 
+clean_test_acceptance: export COMPOSE_PROJECT_NAME=acceptance_test_$(BUILD_DIR_NAME)_$(MODULE_NAME)
 clean_test_acceptance:
-	${DOCKER_COMPOSE_TEST_ACCEPTANCE} down -v -t 0
+	${DOCKER_COMPOSE} down -v -t 0
 
 endif

services/web/app/src/infrastructure/Mongoose.js

@@ -4,15 +4,6 @@ const Metrics = require('@overleaf/metrics')
 const logger = require('@overleaf/logger')
 const { addConnectionDrainer } = require('./GracefulShutdown')
 
-if (
-  typeof global.beforeEach === 'function' &&
-  process.argv.join(' ').match(/unit/)
-) {
-  throw new Error(
-    'It looks like unit tests are running, but you are connecting to Mongo. Missing a stub?'
-  )
-}
-
 mongoose.set('autoIndex', false)
 mongoose.set('strictQuery', false)
 

services/web/app/src/infrastructure/mongodb.js

@@ -16,15 +16,6 @@ if (Mongoose.mongo.ObjectId !== mongodb.ObjectId) {
 
 const { ObjectId, ReadPreference } = mongodb
 
-if (
-  typeof global.beforeEach === 'function' &&
-  process.argv.join(' ').match(/unit/)
-) {
-  throw new Error(
-    'It looks like unit tests are running, but you are connecting to Mongo. Missing a stub?'
-  )
-}
-
 const READ_PREFERENCE_PRIMARY = ReadPreference.primary.mode
 const READ_PREFERENCE_SECONDARY = Settings.mongo.hasSecondaries
   ? ReadPreference.secondary.mode
@@ -101,7 +92,24 @@ async function getCollectionNames() {
   return collections.map(collection => collection.collectionName)
 }
 
+async function cleanupTestDatabase() {
+  ensureTestDatabase()
+  const collectionNames = await getCollectionNames()
+  const collections = []
+  for (const name of collectionNames) {
+    if (name in db && name !== 'migrations') {
+      collections.push(db[name])
+    }
+  }
+  await Promise.all(collections.map(coll => coll.deleteMany({})))
+}
+
 async function dropTestDatabase() {
+  ensureTestDatabase()
+  await mongoClient.db().dropDatabase()
+}
+
+function ensureTestDatabase() {
   const internalDb = mongoClient.db()
   const dbName = internalDb.databaseName
   const env = process.env.NODE_ENV
@@ -111,8 +119,6 @@ async function dropTestDatabase() {
       `Refusing to clear database '${dbName}' in environment '${env}'`
     )
   }
-
-  await internalDb.dropDatabase()
 }
 
 /**
@@ -129,6 +135,7 @@ module.exports = {
   connectionPromise,
   getCollectionNames,
   getCollectionInternal,
+  cleanupTestDatabase,
   dropTestDatabase,
   READ_PREFERENCE_PRIMARY,
   READ_PREFERENCE_SECONDARY,

services/web/docker-compose.ci.yml

@@ -13,10 +13,13 @@ services:
     user: node
     command: npm run test:unit:app
     working_dir: /overleaf/services/web
+    env_file: docker-compose.common.env
     environment:
       BASE_CONFIG:
       OVERLEAF_CONFIG:
       NODE_OPTIONS: "--unhandled-rejections=strict"
+    depends_on:
+      - mongo
 
   test_acceptance:
     build:
@@ -82,6 +85,15 @@ services:
   mongo:
     image: mongo:6.0.13
     command: --replSet overleaf
+    volumes:
+      - ./docker/mongodb-init-replica-set.js:/docker-entrypoint-initdb.d/mongodb-init-replica-set.js
+    environment:
+      MONGO_INITDB_DATABASE: sharelatex
+    extra_hosts:
+      # Required when using the automatic database setup for initializing the
+      # replica set. This override is not needed when running the setup after
+      # starting up mongo.
+      - mongo:127.0.0.1
 
   ldap:
     image: rroemhild/test-openldap:1.1

services/web/docker-compose.yml

@@ -12,6 +12,7 @@ services:
       - ../../node_modules:/overleaf/node_modules
       - ../../libraries:/overleaf/libraries
     working_dir: /overleaf/services/web
+    env_file: docker-compose.common.env
     environment:
       BASE_CONFIG:
       OVERLEAF_CONFIG:
@@ -19,6 +20,8 @@ services:
       NODE_OPTIONS: "--unhandled-rejections=strict"
     command: npm run --silent test:unit:app
     user: node
+    depends_on:
+      - mongo
 
   test_acceptance:
     image: node:20.18.0
@@ -83,6 +86,15 @@ services:
   mongo:
     image: mongo:6.0.13
     command: --replSet overleaf
+    volumes:
+      - ./docker/mongodb-init-replica-set.js:/docker-entrypoint-initdb.d/mongodb-init-replica-set.js
+    environment:
+      MONGO_INITDB_DATABASE: sharelatex
+    extra_hosts:
+      # Required when using the automatic database setup for initializing the
+      # replica set. This override is not needed when running the setup after
+      # starting up mongo.
+      - mongo:127.0.0.1
 
   ldap:
     image: rroemhild/test-openldap:1.1

services/web/docker/mongodb-init-replica-set.js

@@ -0,0 +1,3 @@
+/* eslint-disable no-undef */
+
+rs.initiate({ _id: 'overleaf', members: [{ _id: 0, host: 'mongo:27017' }] })

services/web/test/acceptance/src/helpers/MongoHelper.mjs

@@ -1,7 +1,7 @@
 import { execFile } from 'child_process'
 import {
   connectionPromise,
-  db,
+  cleanupTestDatabase,
   dropTestDatabase,
 } from '../../../../app/src/infrastructure/mongodb.js'
 import Settings from '@overleaf/settings'
@@ -34,16 +34,6 @@ export default {
       })
     })
 
-    afterEach('purge mongo data', async function () {
-      return Promise.all(
-        Object.values(db).map(async collection => {
-          if (collection === db.migrations) {
-            // Do not clear the collection for tracking migrations.
-            return
-          }
-          return collection.deleteMany({})
-        })
-      )
-    })
+    afterEach('purge mongo data', cleanupTestDatabase)
   },
 }

services/web/test/unit/src/Security/OneTimeTokenHandlerTests.js

@@ -1,225 +1,151 @@
-const SandboxedModule = require('sandboxed-module')
-const path = require('path')
 const sinon = require('sinon')
-const modulePath = path.join(
-  __dirname,
-  '../../../../app/src/Features/Security/OneTimeTokenHandler'
-)
-const Errors = require('../../../../app/src/Features/Errors/Errors')
-const tk = require('timekeeper')
 const { expect } = require('chai')
+const Errors = require('../../../../app/src/Features/Errors/Errors')
+const {
+  connectionPromise,
+  cleanupTestDatabase,
+} = require('../../../../app/src/infrastructure/mongodb')
+const OneTimeTokenHandler = require('../../../../app/src/Features/Security/OneTimeTokenHandler')
 
 describe('OneTimeTokenHandler', function () {
+  before(async function () {
+    await connectionPromise
+  })
+  beforeEach(cleanupTestDatabase)
+
   beforeEach(function () {
-    tk.freeze(Date.now()) // freeze the time for these tests
-    this.stubbedToken = 'mock-token'
-    this.callback = sinon.stub()
-    this.OneTimeTokenHandler = SandboxedModule.require(modulePath, {
-      requires: {
-        '@overleaf/settings': this.settings,
-        crypto: {
-          randomBytes: () => this.stubbedToken,
-        },
-        '../../infrastructure/mongodb': {
-          db: (this.db = { tokens: {} }),
-        },
-      },
-    })
+    this.clock = sinon.useFakeTimers()
   })
 
   afterEach(function () {
-    tk.reset()
+    this.clock.restore()
   })
 
   describe('getNewToken', function () {
-    beforeEach(function () {
-      this.db.tokens.insertOne = sinon.stub().yields()
+    it('generates a token and stores it in the database', async function () {
+      const token = await OneTimeTokenHandler.promises.getNewToken(
+        'password',
+        'mock-data-to-store'
+      )
+      const { data } = await OneTimeTokenHandler.promises.peekValueFromToken(
+        'password',
+        token
+      )
+      expect(data).to.equal('mock-data-to-store')
     })
 
-    describe('normally', function () {
-      beforeEach(function () {
-        this.OneTimeTokenHandler.getNewToken(
-          'password',
-          'mock-data-to-store',
-          this.callback
-        )
-      })
-
-      it('should insert a generated token with a 1 hour expiry', function () {
-        this.db.tokens.insertOne
-          .calledWith({
-            use: 'password',
-            token: this.stubbedToken,
-            createdAt: new Date(),
-            expiresAt: new Date(Date.now() + 60 * 60 * 1000),
-            data: 'mock-data-to-store',
-          })
-          .should.equal(true)
-      })
-
-      it('should call the callback with the token', function () {
-        this.callback.calledWith(null, this.stubbedToken).should.equal(true)
-      })
+    it('expires the generated token after 1 hour', async function () {
+      const token = await OneTimeTokenHandler.promises.getNewToken(
+        'password',
+        'mock-data-to-store'
+      )
+      this.clock.tick('25:00:00')
+      await expect(
+        OneTimeTokenHandler.promises.peekValueFromToken('password', token)
+      ).to.be.rejectedWith(Errors.NotFoundError)
     })
 
-    describe('with an optional expiresIn parameter', function () {
-      beforeEach(function () {
-        this.OneTimeTokenHandler.getNewToken(
-          'password',
-          'mock-data-to-store',
-          { expiresIn: 42 },
-          this.callback
-        )
-      })
-
-      it('should insert a generated token with a custom expiry', function () {
-        this.db.tokens.insertOne
-          .calledWith({
-            use: 'password',
-            token: this.stubbedToken,
-            createdAt: new Date(),
-            expiresAt: new Date(Date.now() + 42 * 1000),
-            data: 'mock-data-to-store',
-          })
-          .should.equal(true)
-      })
-
-      it('should call the callback with the token', function () {
-        this.callback.calledWith(null, this.stubbedToken).should.equal(true)
-      })
+    it('accepts an expiresIn parameter', async function () {
+      const token = await OneTimeTokenHandler.promises.getNewToken(
+        'password',
+        'mock-data-to-store',
+        { expiresIn: 42 }
+      )
+      this.clock.tick('00:30')
+      const { data } = await OneTimeTokenHandler.promises.peekValueFromToken(
+        'password',
+        token
+      )
+      expect(data).to.equal('mock-data-to-store')
+      this.clock.tick('00:15')
+      await expect(
+        OneTimeTokenHandler.promises.peekValueFromToken('password', token)
+      ).to.be.rejectedWith(Errors.NotFoundError)
     })
   })
 
   describe('peekValueFromToken', function () {
-    describe('successfully', function () {
+    it('should return the data and peek count', async function () {
       const data = { email: 'some-mock-data' }
-      let result
-      beforeEach(async function () {
-        this.db.tokens.findOneAndUpdate = sinon
-          .stub()
-          .resolves({ data, peekCount: 1 })
-        result = await this.OneTimeTokenHandler.promises.peekValueFromToken(
-          'password',
-          'mock-token'
-        )
-      })
-
-      it('should increment the peekCount', function () {
-        this.db.tokens.findOneAndUpdate
-          .calledWith(
-            {
-              use: 'password',
-              token: 'mock-token',
-              expiresAt: { $gt: new Date() },
-              usedAt: { $exists: false },
-              peekCount: { $not: { $gte: this.OneTimeTokenHandler.MAX_PEEKS } },
-            },
-            {
-              $inc: { peekCount: 1 },
-            }
-          )
-          .should.equal(true)
-      })
-
-      it('should return the data', function () {
-        expect(result).to.deep.equal({ data, remainingPeeks: 3 })
+      const token = await OneTimeTokenHandler.promises.getNewToken(
+        'password',
+        data
+      )
+      const result = await OneTimeTokenHandler.promises.peekValueFromToken(
+        'password',
+        token
+      )
+      expect(result).to.deep.equal({
+        data,
+        remainingPeeks: OneTimeTokenHandler.MAX_PEEKS - 1,
       })
     })
 
-    describe('when a valid token is not found', function () {
-      beforeEach(function () {
-        this.db.tokens.findOneAndUpdate = sinon.stub().resolves(null)
-      })
+    it('should throw a NotFoundError if the token is not found', async function () {
+      await expect(
+        OneTimeTokenHandler.promises.peekValueFromToken('password', 'bad-token')
+      ).to.be.rejectedWith(Errors.NotFoundError)
+    })
 
-      it('should return a NotFoundError', async function () {
-        await expect(
-          this.OneTimeTokenHandler.promises.peekValueFromToken(
-            'password',
-            'mock-token'
-          )
-        ).to.be.rejectedWith(Errors.NotFoundError)
-      })
+    it('should stop returning the data after the peek count is exceeded', async function () {
+      const data = { email: 'some-mock-data' }
+      const token = await OneTimeTokenHandler.promises.getNewToken(
+        'password',
+        data
+      )
+      for (let peeks = 1; peeks <= OneTimeTokenHandler.MAX_PEEKS; peeks++) {
+        const result = await OneTimeTokenHandler.promises.peekValueFromToken(
+          'password',
+          token
+        )
+        expect(result).to.deep.equal({
+          data,
+          remainingPeeks: OneTimeTokenHandler.MAX_PEEKS - peeks,
+        })
+      }
+      await expect(
+        OneTimeTokenHandler.promises.peekValueFromToken('password', token)
+      ).to.be.rejectedWith(Errors.NotFoundError)
     })
   })
 
   describe('expireToken', function () {
-    beforeEach(function () {
-      this.db.tokens.updateOne = sinon.stub().yields(null)
-      this.OneTimeTokenHandler.expireToken(
+    it('should expire the token immediately', async function () {
+      const token = await OneTimeTokenHandler.promises.getNewToken(
         'password',
-        'mock-token',
-        this.callback
+        'mock-data'
       )
-    })
-
-    it('should expire the token', function () {
-      this.db.tokens.updateOne
-        .calledWith(
-          {
-            use: 'password',
-            token: 'mock-token',
-          },
-          {
-            $set: {
-              usedAt: new Date(),
-            },
-          }
-        )
-        .should.equal(true)
-      this.callback.calledWith(null).should.equal(true)
+      await OneTimeTokenHandler.promises.expireToken('password', token)
+      await expect(
+        OneTimeTokenHandler.promises.peekValueFromToken('password', token)
+      ).to.be.rejectedWith(Errors.NotFoundError)
     })
   })
 
   describe('getValueFromTokenAndExpire', function () {
-    describe('successfully', function () {
-      beforeEach(function () {
-        this.db.tokens.findOneAndUpdate = sinon
-          .stub()
-          .yields(null, { data: 'mock-data' })
-        this.OneTimeTokenHandler.getValueFromTokenAndExpire(
+    it('should return the value and expire the token', async function () {
+      const token = await OneTimeTokenHandler.promises.getNewToken(
+        'password',
+        'mock-data'
+      )
+      const data =
+        await OneTimeTokenHandler.promises.getValueFromTokenAndExpire(
           'password',
-          'mock-token',
-          this.callback
+          token
         )
-      })
-
-      it('should expire the token', function () {
-        this.db.tokens.findOneAndUpdate
-          .calledWith(
-            {
-              use: 'password',
-              token: 'mock-token',
-              expiresAt: { $gt: new Date() },
-              usedAt: { $exists: false },
-              peekCount: { $not: { $gte: this.OneTimeTokenHandler.MAX_PEEKS } },
-            },
-            {
-              $set: { usedAt: new Date() },
-            }
-          )
-          .should.equal(true)
-      })
-
-      it('should return the data', function () {
-        this.callback.calledWith(null, 'mock-data').should.equal(true)
-      })
+      expect(data).to.equal('mock-data')
+      await expect(
+        OneTimeTokenHandler.promises.peekValueFromToken('password', token)
+      ).to.be.rejectedWith(Errors.NotFoundError)
     })
 
-    describe('when a valid token is not found', function () {
-      beforeEach(function () {
-        this.db.tokens.findOneAndUpdate = sinon.stub().yields(null, null)
-        this.OneTimeTokenHandler.getValueFromTokenAndExpire(
+    it('should throw a NotFoundError if the token is not found', async function () {
+      await expect(
+        OneTimeTokenHandler.promises.getValueFromTokenAndExpire(
           'password',
-          'mock-token',
-          this.callback
+          'bad-token'
         )
-      })
-
-      it('should return a NotFoundError', function () {
-        this.callback
-          .calledWith(sinon.match.instanceOf(Errors.NotFoundError))
-          .should.equal(true)
-      })
+      ).to.be.rejectedWith(Errors.NotFoundError)
     })
   })
 })