Post test cleanup. Closes #577

README.md

@@ -183,6 +183,24 @@ lab.experiment('with only', () => {
 });
 ```
 
+The `test()` callback provides a second `onCleanup` argument which is a function used to register a runtime cleanup function
+to be executed after the test completed. The cleanup function will execute even in the event of a timeout. Note that the cleanup
+function will be executed as-is without any timers and if it fails to call it's `next` argument, the runner will freeze.
+
+```javascript
+lab.test('cleanups after test', (done, onCleanup) => {
+
+    onCleanup((next) => {
+
+        cleanup_logic();
+        return next();
+    });
+
+    Code.expect(1 + 1).to.equal(2);
+    done();
+});
+```
+
 Additionally, `test()` options support a `plan` setting to specify the expected number of assertions for your test to execute. This
 setting should only be used with an assertion library that supports a `count()` function, like [`code`](http://npmjs.com/package/code).
 *`plan` may not work with parallel test executions*

lib/runner.js

@@ -537,7 +537,11 @@ internals.protect = function (item, state, callback) {
 
         const immed = setImmediate(() => {
 
-            return callback(err);
+            if (!item.onCleanup) {
+                return callback(err);
+            }
+
+            item.onCleanup(() => callback(err));
         });
 
         /* $lab:coverage:off$ */
@@ -566,7 +570,10 @@ internals.protect = function (item, state, callback) {
 
         // 1. Do not forward what's already a forward.
         // 2. Only errors that reach before*/after* are worth forwarding, otherwise we know where they came from.
-        if (!isForward && item.id === undefined) {
+
+        if (!isForward &&
+            item.id === undefined) {
+
             internals.forwardError(err, domain, domains);
         }
 
@@ -584,16 +591,30 @@ internals.protect = function (item, state, callback) {
     setImmediate(() => {
 
         domain.enter();
+
+        item.onCleanup = null;
+        const onCleanup = (func) => {
+
+            item.onCleanup = func;
+        };
+
         const itemResult = item.fn.call(null, (err) => {
 
             finish(err, 'done');
-        });
-        if (itemResult && itemResult.then instanceof Function) {
+        }, onCleanup);
+
+        if (itemResult &&
+            itemResult.then instanceof Function) {
+
             itemResult.then(() => finish(null), (err) => finish(err, 'done'));
         }
-        if (item.fn.length !== 1 && !(itemResult && itemResult.then instanceof Function)) {
+
+        if (item.fn.length !== 1 &&
+            !(itemResult && itemResult.then instanceof Function)) {
+
             finish(new Error(`Function for "${item.title}" should either take a callback argument or return a promise`), 'function signature');
         }
+
         domain.exit();
     });
 };

test/runner.js

@@ -99,6 +99,32 @@ describe('Runner', () => {
         });
     });
 
+    it('calls cleanup function', (done) => {
+
+        const script = Lab.script();
+
+        let flag = false;
+        script.test('a', (done, onCleanup) => {
+
+            onCleanup((next) => {
+
+                flag = true;
+                return next();
+            });
+
+            done();
+        });
+
+        Lab.execute(script, {}, null, (err, notebook) => {
+
+            expect(err).not.to.exist();
+            expect(notebook.tests).to.have.length(1);
+            expect(notebook.failures).to.equal(0);
+            expect(flag).to.be.true();
+            done();
+        });
+    });
+
     it('should fail test that neither takes a callback nor returns anything', (done) => {
 
         const script = Lab.script({ schedule: false });