From 63822c5124ddfcfb6a77df6c98dcedd37c9885c1 Mon Sep 17 00:00:00 2001
From: "Harish.K" <harish2704@gmail.com>
Date: Fri, 1 Apr 2016 01:32:55 +0530
Subject: [PATCH] Fix: make all functions as async Feat: afterFn parsing

---
 index-v1.js | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/index-v1.js b/index-v1.js
index 934a7d2..2194ede 100644
--- a/index-v1.js
+++ b/index-v1.js
@@ -7,6 +7,9 @@ var fs = require('fs');
 
 
 function mkAsyncFn( fn ){
+  if( fn.length ){
+    return fn;
+  }
   return function( done ){
     setTimeout( function(){
       fn();
@@ -18,10 +21,8 @@ function mkAsyncFn( fn ){
 
 
 function ItBlock( description, fn ){
-  var isAsync = ( fn.length > 0 );
-
   this.description = description;
-  this.fn = isAsync ? fn : mkAsyncFn( fn );
+  this.fn = mkAsyncFn( fn );
 }
 
 
@@ -84,11 +85,13 @@ SimpleMocha.prototype.it         = function( description, fn ){
 
 
 SimpleMocha.prototype.before     = function( fn ){
-  this.currentDescribeBlock.beforeFn = fn;
+  this.currentDescribeBlock.beforeFn = mkAsyncFn( fn );
 };
 
 
-SimpleMocha.prototype.after      = function(){};
+SimpleMocha.prototype.after      = function( fn ){
+  this.currentDescribeBlock.afterFn = mkAsyncFn( fn );
+};
 SimpleMocha.prototype.beforeEach = function(){};
 SimpleMocha.prototype.afterEach  = function(){};