Mockingoose Vitest

mockingvite.ts

/*
 * A Vitest compatible version of Mockingoose
 * Adapted from https://gist.github.com/stolinski/9b1de328b79e7b6edd8def14c224afdf
 */

import mongoose from 'mongoose';
import { vi } from 'vitest';

mongoose.connect = vi.fn().mockImplementation(() => Promise.resolve());

mongoose.createConnection = vi.fn().mockReturnValue({
  catch() {
    /* no op */
  },
  model: mongoose.model.bind(mongoose),
  on: vi.fn(),
  once: vi.fn(),
  then(resolve) {
    return Promise.resolve(resolve(this));
  },
});

const ops = [
  'find',
  'findOne',
  'count',
  'countDocuments',
  'estimatedDocumentCount',
  'distinct',
  'findOneAndUpdate',
  'findOneAndDelete',
  'findOneAndRemove',
  'findOneAndReplace',
  'remove',
  'update',
  'updateOne',
  'updateMany',
  'deleteOne',
  'deleteMany',
  'save',
  'aggregate',
  '$save',
];

const mockedReturn = async function (cb) {
  const {
    op,
    model: { modelName },
    _mongooseOptions = {},
  } = this;
  const Model = mongoose.model(modelName);

  let mock =
    mockingoose.__mocks[modelName] && mockingoose.__mocks[modelName][op];

  let err = null;

  if (mock instanceof Error) {
    err = mock;
  }

  if (typeof mock === 'function') {
    mock = await mock(this);
  }

  if (!mock && op === 'save') {
    mock = this;
  }

  if (!mock && op === '$save') {
    mock = this;
  }

  if (
    mock &&
    !(mock instanceof Model) &&
    ![
      'remove',
      'deleteOne',
      'deleteMany',
      'update',
      'updateOne',
      'updateMany',
      'count',
      'countDocuments',
      'estimatedDocumentCount',
      'distinct',
    ].includes(op)
  ) {
    mock = Array.isArray(mock)
      ? mock.map((item) => new Model(item))
      : new Model(mock);

    if (op === 'insertMany') {
      if (!Array.isArray(mock)) mock = [mock];

      for (const doc of mock) {
        const e = doc.validateSync();
        if (e) throw e;
      }
    }

    if (_mongooseOptions.lean || _mongooseOptions.rawResult) {
      mock = Array.isArray(mock)
        ? mock.map((item) => item.toObject())
        : mock.toObject();
    }
  }

  if (cb) {
    return cb(err, mock);
  }

  if (err) {
    throw err;
  }

  return mock;
};

ops.forEach((op) => {
  mongoose.Query.prototype[op] = vi
    .fn()
    .mockImplementation(function (criteria, doc, options, callback) {
      // Why is this here?
      let crit = criteria;
      let ddoc = doc;
      let opts = options;
      let cb = callback;
      switch (arguments.length) {
        case 4:
        case 3:
          if (typeof options === 'function') {
            cb = options;
            opts = {};
          }
          break;
        case 2:
          if (typeof doc === 'function') {
            cb = doc;
            ddoc = criteria;
            crit = undefined;
          }
          opts = undefined;
          break;
        case 1:
          if (typeof criteria === 'function') {
            cb = criteria;
            crit = opts = ddoc = undefined;
          } else {
            ddoc = criteria;
            crit = opts = undefined;
          }
      }

      if (
        [
          'find',
          'findOne',
          'count',
          'countDocuments',
          'remove',
          'deleteOne',
          'deleteMany',
          'update',
          'updateOne',
          'updateMany',
          'findOneAndUpdate',
          'findOneAndRemove',
          'findOneAndDelete',
          'findOneAndReplace',
        ].includes(op) &&
        typeof crit !== 'function'
      ) {
        // find and findOne can take conditions as the first paramter
        // ensure they make it into the Query conditions
        this.merge(crit);
      }

      if (['distinct'].includes(op) && typeof ddoc !== 'function') {
        // distinct has the conditions as the second parameter
        this.merge(ddoc);
      }

      if (/update/i.test(op) && typeof ddoc !== 'function' && ddoc) {
        this.setUpdate(ddoc);
      }

      this.op = op;

      if (!cb) {
        return this;
      }

      return this.exec.call(this, cb);
    });
});

mongoose.Query.prototype.exec = vi.fn().mockImplementation(function (cb) {
  return mockedReturn.call(this, cb);
});

mongoose.Aggregate.prototype.exec = vi
  .fn()
  .mockImplementation(async function (cb) {
    const {
      _model: { modelName },
    } = this;

    let mock =
      mockingoose.__mocks[modelName] &&
      mockingoose.__mocks[modelName].aggregate;

    let err = null;

    if (mock instanceof Error) {
      err = mock;
    }

    if (typeof mock === 'function') {
      mock = await mock(this);
    }

    if (cb) {
      return cb(err, mock);
    }

    if (err) {
      throw err;
    }

    return mock;
  });

mongoose.Model.insertMany = vi
  .fn()
  .mockImplementation(function (arr, options, cb) {
    const op = 'insertMany';
    const { modelName } = this;

    let actualCb = cb;
    if (typeof options === 'function') {
      actualCb = options;
    } else {
      this._mongooseOptions = options;
    }

    Object.assign(this, { model: { modelName }, op });
    return mockedReturn.call(this, actualCb);
  });

const instance = ['remove', 'save', '$save'];

instance.forEach((methodName) => {
  mongoose.Model.prototype[methodName] = vi
    .fn()
    .mockImplementation(function (options, cb) {
      const op = methodName;
      const { modelName } = this.constructor;

      let actualCb = cb;
      if (typeof options === 'function') {
        actualCb = options;
      }

      Object.assign(this, { model: { modelName }, op });

      const hooks = this.constructor.hooks;

      return new Promise((resolve, reject) => {
        hooks.execPre(op, this, [actualCb], (err) => {
          if (err) {
            reject(err);
            return;
          }

          const ret = mockedReturn.call(this, actualCb);

          if (actualCb) {
            hooks.execPost(op, this, [ret], (err2) => {
              if (err2) {
                reject(err2);
                return;
              }

              resolve(ret);
            });
          } else {
            ret
              .then((ret2) => {
                hooks.execPost(op, this, [ret2], (err3) => {
                  if (err3) {
                    reject(err3);
                    return;
                  }

                  resolve(ret2);
                });
              })
              .catch(reject);
          }
        });
      });
    });
});

vi.doMock('mongoose', () => mongoose);

// extend a plain function, we will override it with the Proxy later
const proxyTarget = Object.assign(() => void 0, {
  __mocks: {},
  resetAll() {
    this.__mocks = {};
  },
  toJSON() {
    return this.__mocks;
  },
});

const getMockController = (prop) => {
  return {
    reset(op) {
      if (op) {
        delete proxyTarget.__mocks[prop][op];
      } else {
        delete proxyTarget.__mocks[prop];
      }

      return this;
    },
    toJSON() {
      return proxyTarget.__mocks[prop] || {};
    },
    toReturn(o, op = 'find') {
      Object.prototype.hasOwnProperty.call(proxyTarget.__mocks, prop)
        ? (proxyTarget.__mocks[prop][op] = o)
        : (proxyTarget.__mocks[prop] = { [op]: o });

      return this;
    },
  };
};

const proxyTraps = {
  apply: (target, thisArg, [prop]) => mockModel(prop),
  get(target, prop) {
    if (Object.prototype.hasOwnProperty.call(target, prop)) {
      return Reflect.get(target, prop);
    }

    return getMockController(prop);
  },
};

const mockingoose = new Proxy(proxyTarget, proxyTraps);

/**
 * Returns a helper with which you can set up mocks for a particular Model
 */
const mockModel = (model) => {
  const modelName = typeof model === 'function' ? model.modelName : model;
  if (typeof modelName === 'string') {
    return getMockController(modelName);
  } else {
    throw new Error('model must be a string or mongoose.Model');
  }
};

export default mockingoose;