Labnotes

Yield to the test: using Mocha with ES6 generators

Published on

Why would I want to?

You can yield with abandon:

describe("New customer", function() {
  var business;
  var customer;

  before(function*() {
    yield setup();
    business = yield Account.create("FooBar Inc");
    customer = yield business.addCustomer("Mr. Baz");
  });
  
  it("should be the only customer", function*() {
    var count = yield Customer.count({ businessID: business.id });
    assert.equal(count, 1);
  });
  
  after(function*() {
    yield teardown();
  });

});

Silly made up test case, but you get the point: you can use generators instead of callbacks in before hooks, after hooks, and of course, test cases.

If you're testing with Zombie.js, you can do this:

before(function*() {
  yield browser.visit('/signup/new');
  browser.fill('name', 'Assaf');
  browser.fill('password', 'Secret');
  yield browser.pressButton('Signup');
});

So let's make it happen.

1. Add ES6 Generators

This depends on your environment. You can run Node with the --harmony flag. If you're using Traceur, you can enable ES6 by running this first:

var Traceur = require('traceur');

// Traceur will compile all JS aside from node modules
Traceur.require.makeDefault(function(filename) {
  return !(/node_modules/.test(filename));
});

2. Enhance Mocha

Fortunately, Mocha runs everything using the Runnable class, all we need to do is add generator support to the run method.

This is mostly Mocha code, with one addition: we're looking for synchronous functions (that take no arguments), and if the return a generator, we use co to run it.

I picked co because it understands promises, and can run code in parallel. suspend and genny should work just as well.

const co    = require('co');
const mocha = require('mocha');


mocha.Runnable.prototype.run = function(fn) {
  var self = this
    , ms = this.timeout()
    , start = new Date
    , ctx = this.ctx
    , finished
    , emitted;

  if (ctx) ctx.runnable(this);

  // timeout
  if (this.async) {
    if (ms) {
      this.timer = setTimeout(function(){
        done(new Error('timeout of ' + ms + 'ms exceeded'));
        self.timedOut = true;
      }, ms);
    }
  }

  // called multiple times
  function multiple(err) {
    if (emitted) return;
    emitted = true;
    self.emit('error', err || new Error('done() called multiple times'));
  }

  // finished
  function done(err) {
    if (self.timedOut) return;
    if (finished) return multiple(err);
    self.clearTimeout();
    self.duration = new Date - start;
    finished = true;
    fn(err);
  }

  // for .resetTimeout()
  this.callback = done;

  // async
  if (this.async) {
    try {
      this.fn.call(ctx, function(err){
        if (err instanceof Error || toString.call(err) === "[object Error]") return done(err);
        if (null != err) return done(new Error('done() invoked with non-Error: ' + err));
        done();
      });
    } catch (err) {
      done(err);
    }
    return;
  }

  if (this.asyncOnly) {
    return done(new Error('--async-only option in use without declaring `done()`'));
  }

  try {
    if (!this.pending) {
      var result = this.fn.call(ctx);
      // This is where we determine if the result is a generator
      if (result && typeof(result.next) == 'function' && typeof(result.throw) == 'function') {
        // Mocha timeout for async function
        if (ms) {
          this.timer = setTimeout(function(){
            done(new Error('timeout of ' + ms + 'ms exceeded'));
            self.timedOut = true;
          }, ms);
        }
        // Use co to run generator to completion
        co(result)(function(err) {
          this.duration = new Date - start;
          done(err);
        });
      } else {
        // Default Mocha handling of sync function
        this.duration = new Date - start;
        fn();
      }
    }
  } catch (err) {
    fn(err);
  }

}

3. Run at start

We're going to take these enhancements and place them in a file called test/es6_mocha.js.

Next, we need to tell Mocha to run this file before loading any test cases, and we do that by adding the following line to test/mocha.opts:

--require test/es6_mocha.js

4. You're done

Enjoy!