Simple mocking library with JavaScript proxy

anakin - darth vader duality

One of the coolest features in JS is a Proxy object. Pity I can’t find more uses for them in a day-to-day development. Or maybe that’s a good thing.

While there are mocking libraries out there, this is a reinvent-the-wheel approach, purely for fun. Because why not…

In a nutshell, proxies allow us to intercept if an object property getter or setter was invoked, and then tamper with it. Here’s a simple example:

const lib = {
    version: '1.0',
    square (x: number) { return x * x; }
}

Proxy objects are simple to make. The constructor takes an object we wish to tamper with, and a set of getter/setter “traps”.

We’ll get to function interception later. For now, let’s focus on version property. We can write a simple proxy for this lib object:

const libProxy = new Proxy(lib, {
    get (target, prop) {
        if (prop === 'version') {
            return target[prop] + '-beta';
        }
    },
    set(target, prop, value) {
        if (prop === 'version') {
            target[prop] = '1.0';
        }

        return true;
    }
});

The get trap has two arguments here. target is the object we’re proxying, and prop is a property for which getter was invoked. If you call:

console.log(libProxy.version) // '1.0-beta'

as specified above, it will attach -beta suffix to version value. set function has an additional argument, and that’s a value someone might try to set. In this example, no matter what you try, version will always remain 1.0 under the hood, but it will return 1.0-beta ‘cause we’ve tampered with a getter.

libProxy.version = 'v25.17.6-prealphabuildtotallystable'
console.log(libProxy.version) // '1.0-beta'

Functions are no different. The example below just shows set trap, but you get the picture:

const libProxy = new Proxy(lib, {
    set(target, prop, value) {
        if (prop === 'square') {
            target[prop] = value;
        }

        return true;
    }
});

libProxy.square = (x: number) => x * x * x * x; // fourth power
console.log(libProxy.square(2)); // 16

Considering the title of this article, you should realize where this is going. We already have a crude mocking functionality. Trouble is, you’d have to create this proxy over and over and over again for each test. It would be better to just have something like this:

mock.someFn.whenCalledWith('blah').returns('derp');

And this is what we’re heading towards.

Automation

This is gonna require some advanced type juggling, if we want to do it in TS. So let’s start with the first help type:

type AnyFn = (...args: any[]) => any;

It defines a reusable function type alias for any function signature that takes any argument(s) and returns any value.

type MockedFn<F extends AnyFn> = ((...args: Parameters<F>) => ReturnType<F>) & {
  calls: Parameters<F>[];
  setImpl: (impl: F) => void;
  reset: () => void;
  whenCalledWith: (
    ...args: Parameters<F>
  ) => {
    returns: (value: ReturnType<F>) => void;
    callsFake: (impl: (...a: Parameters<F>) => ReturnType<F>) => void;
  };
};

This is a generic mock-function type that wraps a real function type Fand keeps its argument and return types. Parameters<> and ReturnType are build-in TS types.

calls property stores every argument list the mock was called with.

setImpl method is used to set the mock’s fallback implementation.

Get Bananica Bananica’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

reset clears calls and configured behavior.

whenCalledWith takes same parameters as F (the real function we’re trying to mock) and returns an object with 2 functions. returns function just sets what value to return while callsFake allows to execute random function body with same parameters (maybe we wanna do some side effects, but that’s usually not recommended). This would eventually allow us to write what was mentioned above:

mock.someFn.whenCalledWith('blah').returns('derp');
// or
mock.someFn.whenCalledWith('blah').callsFake(x => x + ' blah'); // 'blah blah'

Third and final helper type:

type Mocked<T extends object> = {
  [K in keyof T]: T[K] extends AnyFn ? MockedFn<T[K]> : T[K];
};

Only objects can be mocked. This type creates a mapping of the original object. For each key of the original, convert function properties, and cast them into MockedFn. Leave regular properties as they are.

Now it’s time to implement createMock function. Let’s go over signature first:

function createMock<T extends object>(
  defaults: Partial<{ [K in keyof T]: T[K] }> = {},
): Mocked<T> {
  // ...implementation
}

Again T must extend an object. defaults is an object of pre-supplied values for the mock. { [K in keyof T]: T[K] } recreates the shape of T. Partial<...> makes everything optional, so only properties that we want can be passed. {} is an empty default parameter object, so we can get mock just by declaring type:

const mock = createMock<Foo>();

… oh, and it returns Mocked<T> of course. Moving on:

function createMock<T extends object>(
  defaults: Partial<{ [K in keyof T]: T[K] }> = {},
): Mocked<T> {
  const store = new Map<PropertyKey, any>();

  const sameArgs = (a: unknown[], b: unknown[]) =>
    a.length === b.length && a.every((x, i) => Object.is(x, b[i]));

  const proxy = new Proxy({} as Mocked<T>, {
    // implemented below
  });
  
  return proxy;
}

store is an internal storage for generated functions/properties.

sameArgs is a helper function which compares passed arguments.

Finally, we need to implement and return the proxy. Setter trap is simple:

const proxy = new Proxy({} as Mocked<T>, {
  set(_target, prop, value) {
    store.set(prop, value);
    return true;
  },
  get() {
    // implemented below
  }
});

And the getter (explainer in the comments, since there’s a lot):

get(_target, prop: keyof T) {
      // if property exists in defaults, just return it
      if (prop in defaults) return (defaults as any)[prop];

      // If property is not in store yet,
      // initialize a new mock function for it.
      if (!store.has(prop)) {
        // array to record all calls to this mock function
        const calls: any[] = [];
        
        // optional fallback implementation function
        let impl: AnyFn | undefined;

        // handlers array type for argument-specific behaviors
        const handlers: Array<{
          // each handler stores the expected argument list
          args: unknown[];
          // each handler stores a resolver function to produce a return value
          resolve: (...args: unknown[]) => unknown;
        }> = [];

        // defines the actual mock function body.
        const fn: AnyFn = (...args: any[]) => {
          calls.push(args);
          // finds first handler whose expected args match actual args
          const hit = handlers.find((h) => sameArgs(h.args, args));
          
          // if matching handler exists, returns its resolved value
          if (hit) return hit.resolve(...args);
          
          // if no handler matched but fallback implementation exists, calls fallback
          if (impl) return impl(...args);
          
          // if nothing is matched, return undefined
          return undefined;
        };

        // add extra properties/methods to the mocked function object
        Object.defineProperties(fn, {
          // defines calls getter so callers can inspect call history
          calls: { get: () => calls },
          // defines setImpl implementation
          setImpl: {
            value: (newImpl: AnyFn) => {
              // assign new fallback implementation
              impl = newImpl;
            },
          },
          // reset mock implementation
          reset: {
            value: () => {
              calls.length = 0;
              impl = undefined;
              handlers.length = 0;
            },
          },
          // define whenCalledWith function returning configuration helpers
          whenCalledWith: {
            value: (...expectedArgs: any[]) => ({
              returns: (value: any) => {
                // add handler that always returns a fixed value for expected args
                handlers.push({ args: expectedArgs, resolve: () => value });
              },
              callsFake: (fakeImpl: AnyFn) => {
                handlers.push({
                  // stores expected args for this handler
                  args: expectedArgs,
                  // resolver calls provided fake implementation with actual args
                  resolve: (...actualArgs: any[]) => fakeImpl(...actualArgs),
                });
              },
            }),
          },
        });

        // store generated mock function in map under this property key
        store.set(prop, fn);
      } // end if(!store.has(prop))

      // get prop from store and return it
      return store.get(prop);
}

Whew… That’s a lot honestly. I understand it probably isn’t easy to swallow, but playing around with this code in the editor will give you a better idea. Now it can be tested:

type Lib = {
  version: string;
  find: (id: number) => { id: number, data: string }
}

const mock = createMock<Lib>({ version: "mock" }); // optional object with version

mock.find.whenCalledWith(1).returns({ id: 1, data: 'junk' });
// or
mock.find.whenCalledWith(1).callsFake(x => ({ id: x, data: 'junk'}));

const result1 = mock.find(1); // { id: 1, data: 'junk' }
const result2 = mock.find(2); // undefined

const mockedCallsArgs = mock.find.calls; // [ [1], [2] ]

mock.find.reset();

You can also pass mocked object as a dependency, no problem:

function doStuff(lib: Lib) {
  const dataSet = [1,2,3].forEach(x => lib.find);
  return dataSet;
}

const mock = createMock<Lib>({ version: "mock" });
mock.find.whenCalledWith(1).returns({ id: 1, data: 'junk' });

doStuff(mock);

Works for async functions too. Give it a spin.

Conclusion

This one has some more advanced concepts, true, and I’m sure someone can find edge cases this implementation doesn’t cover. But honestly, it’s pretty solid for ~100 lines of code.

I hope this was interesting and informative article.

Thank You for reading 🙋‍♂️🙃