wrap.js

const PRESTART = 0
  ,   STARTED = 1
  ,   STARTDONE = 2
  ,   STARTERRORED = 3
  ;
/* convenience function for wrapping a suite with setup and teardown
 * functions. this takes an object which has three properties:
 *
 * (by the way, I'm looking for a better name for this function)
 *
 * suite:          the test suite object, required
 * setup:          a function that should be run before the test
 * teardown:       a function that should be run after the test
 * teardown:       a function that should be run after the test
 * suiteSetup:     a function run before any tests in a suite are ran
 */
exports.wrap = function(obj) {
  var suite = obj.suite
    , setup = obj.setup
    , teardown = obj.teardown
    , suiteSetup = obj.suiteSetup
    ;

  if (!suite) {
    throw new Error('Cannot wrap suite -- no suite provided');
  }

  for(var key in suite) {
    if (typeof suite[key] == 'function') {
      suite[key] = wrapper(suite[key]);
    }
    else {
      exports.wrap({suite: suite[key], setup: setup, teardown: teardown});
    }
  }

  var state = 0
    , pendingTests = []
    , controlError
    ;

  function wrapper(func) {
    var newFunc = function(test) {
      if (teardown) {
        var finish = test.finish;
        test.finish = function() {
          teardown(test, finish);
        }
      }

      if (setup) {
        setup(test, function() { func(test); });
      }
      else {
        func(test);
      }
    }

    if (suiteSetup) {
      // for async suiteStartup
      var setupDone = function() {
        state = STARTDONE;
        process.removeListener('uncaughtException', setupErrored);
        for (var i = 0; i < pendingTests.length; i++) {
          pendingTests[i][0](pendingTests[i][1]);
        }
      }

      var setupErrored = function(err) {
        state = STARTERRORED;
        controlError = err;

        if (pendingTests.length > 1) {
          pendingTests.pop();
          process.removeListener('uncaughtException', setupErrored);
          process.nextTick(function() {
            throw err;
          });
        }
      }

      // wrapping newFunc
      var f = newFunc;
      newFunc = function(test) {
        switch(state) {
          case PRESTART: // suiteSetup hasn't been ran
            state = STARTED;
            process.on('uncaughtException', setupErrored);
            pendingTests.push([f, test]);

            try {
              suiteSetup(setupDone);
            }
            catch(e) {
              setupErrored(e);
            }
            break;
          case STARTED: // suiteSetup is running
            pendingTests.push([f, test]);
            break;
          case STARTDONE: // suiteSetup is done
            f(test);
            break;
          case STARTERRORED: // need to error the test
            (function() {
              throw controlError;
            })();
            break;
        }
      }
    }

    newFunc.toString = function() {
      return  (suiteSetup ? suiteSetup+'\n' : '') +
              (setup ? setup+'\n' : '') +
              func +
              (teardown ? '\n'+teardown : '')
              ;
    }

    return newFunc;
  }

  return suite;
}