From bb0a321b5a0bdf9d326704d2c216d64c25c5cc77 Mon Sep 17 00:00:00 2001
From: Kamal Mahyuddin <kamal.fariz@gmail.com>
Date: Sat, 21 Mar 2015 23:59:02 -0700
Subject: [PATCH] yield assert into the beforeEach/afterEach callbacks if they
 contain assertions, fixes abuiles/ember-watson#4

---
 lib/formulas/qunit-transforms.js             | 63 +++++++++++++++++---
 tests/fixtures/qunit-files/new-with-qunit.js |  3 +-
 tests/fixtures/qunit-files/old-with-qunit.js |  1 +
 3 files changed, 58 insertions(+), 9 deletions(-)

diff --git a/lib/formulas/qunit-transforms.js b/lib/formulas/qunit-transforms.js
index 307150b..9631b62 100644
--- a/lib/formulas/qunit-transforms.js
+++ b/lib/formulas/qunit-transforms.js
@@ -22,6 +22,24 @@ function isModule(node) {
     node.expression.callee.name === 'module';
 }
 
+function isSetup(path) {
+  var node = path.node;
+  var parent = path.parent.node;
+
+  return parent.type === 'Property' &&
+    parent.key.name === 'setup' &&
+    node.type === 'FunctionExpression';
+}
+
+function isTeardown(path) {
+  var node = path.node;
+  var parent = path.parent.node;
+
+  return parent.type === 'Property' &&
+    parent.key.name === 'teardown' &&
+    node.type === 'FunctionExpression';
+}
+
 function isTest(node) {
   return node.expression.type === 'CallExpression' &&
     node.expression.callee.name === 'test';
@@ -42,16 +60,22 @@ function isImportFor(module, node) {
     node.source.value === module;
 }
 
-function transformModule(node) {
+function transformModule(node, withSetupAssert, withTeardownAssert) {
   var callExpression = node.expression;
   if (callExpression.arguments.length > 1 &&
       callExpression.arguments[1].type === 'ObjectExpression') {
     callExpression.arguments[1].properties.forEach(function(node) {
       if (node.key.name === 'setup') {
         node.key.name = 'beforeEach';
+        if (withSetupAssert) {
+          transformTestCallback(node.value);
+        }
       }
       if (node.key.name === 'teardown') {
         node.key.name = 'afterEach';
+        if (withTeardownAssert) {
+          transformTestCallback(node.value);
+        }
       }
     });
   }
@@ -59,12 +83,14 @@ function transformModule(node) {
 
 function transformTestStatement(node) {
   if (node.expression.arguments.length > 1) {
-    var callback = node.expression.arguments[1];
+    transformTestCallback(node.expression.arguments[1]);
+  }
+}
 
-    if (callback.type === 'FunctionExpression') {
-      if (callback.params.length === 0) {
-        callback.params.push(builders.identifier('assert'));
-      }
+function transformTestCallback(callback){
+  if (callback.type === 'FunctionExpression') {
+    if (callback.params.length === 0) {
+      callback.params.push(builders.identifier('assert'));
     }
   }
 }
@@ -98,12 +124,15 @@ module.exports = function transform(source) {
   var sections = {
     addQUnitImport: true,
     modules: [],
+    moduleCallbacks: [],
     tests: [],
     skips: [],
     assertions: []
   };
 
   var ast = recast.parse(source);
+  var isInSetup = false;
+  var isInTeardown = false;
 
   recast.visit(ast, {
     visitExpressionStatement: function(path) {
@@ -111,6 +140,7 @@ module.exports = function transform(source) {
 
       if (isModule(node)) {
         sections.modules.push(node);
+        sections.moduleCallbacks.push({});
       }
 
       if (isTest(node)) {
@@ -123,6 +153,14 @@ module.exports = function transform(source) {
 
       if (isAssertion(node)) {
         sections.assertions.push(node);
+
+        if (isInSetup) {
+          sections.moduleCallbacks[sections.moduleCallbacks.length - 1].withSetupAssert = true;
+        }
+
+        if (isInTeardown) {
+          sections.moduleCallbacks[sections.moduleCallbacks.length - 1].withTeardownAssert = true;
+        }
       }
       this.traverse(path);
     },
@@ -137,11 +175,20 @@ module.exports = function transform(source) {
       }
 
       this.traverse(path);
+    },
+    visitFunctionExpression: function(path) {
+      isInSetup = isSetup(path);
+      isInTeardown = isTeardown(path);
+
+      this.traverse(path);
+
+      isInSetup = false;
+      isInTeardown = false;
     }
   });
 
-  sections.modules.forEach(function(qunitModule) {
-    transformModule(qunitModule);
+  sections.modules.forEach(function(qunitModule, index) {
+    transformModule(qunitModule, sections.moduleCallbacks[index].withSetupAssert, sections.moduleCallbacks[index].withTeardownAssert);
   });
 
   sections.tests.forEach(function(node) {
diff --git a/tests/fixtures/qunit-files/new-with-qunit.js b/tests/fixtures/qunit-files/new-with-qunit.js
index 565335c..9b96eaa 100644
--- a/tests/fixtures/qunit-files/new-with-qunit.js
+++ b/tests/fixtures/qunit-files/new-with-qunit.js
@@ -7,8 +7,9 @@ var App;
 var server;
 
 module('Acceptance: FriendsNew', {
-  beforeEach: function() {
+  beforeEach: function(assert) {
     App = startApp();
+    assert.ok(true, 'app started');
   },
   afterEach: function() {
     Ember.run(App, 'destroy');
diff --git a/tests/fixtures/qunit-files/old-with-qunit.js b/tests/fixtures/qunit-files/old-with-qunit.js
index 675cd60..f8d85cc 100644
--- a/tests/fixtures/qunit-files/old-with-qunit.js
+++ b/tests/fixtures/qunit-files/old-with-qunit.js
@@ -8,6 +8,7 @@ var server;
 module('Acceptance: FriendsNew', {
   setup: function() {
     App = startApp();
+    ok(true, 'app started');
   },
   teardown: function() {
     Ember.run(App, 'destroy');