diff --git a/graphs_trees/check_balance/check_balance_challenge.ipynb b/graphs_trees/check_balance/check_balance_challenge.ipynb
index bcd1fb4..3938bdd 100644
--- a/graphs_trees/check_balance/check_balance_challenge.ipynb
+++ b/graphs_trees/check_balance/check_balance_challenge.ipynb
@@ -36,7 +36,7 @@
     "\n",
     "* Is a balanced tree one where the heights of two sub trees of any node doesn't differ by more than 1?\n",
     "    * Yes\n",
-    "* If this is called on a None input, should we return False?\n",
+    "* If this is called on a None input, should we raise an exception?\n",
     "    * Yes\n",
     "* Can we assume we already have a Node class with an insert method?\n",
     "    * Yes\n",
@@ -121,13 +121,16 @@
    "source": [
     "# %load test_check_balance.py\n",
     "from nose.tools import assert_equal\n",
+    "from nose.tools import raises\n",
     "\n",
     "\n",
     "class TestCheckBalance(object):\n",
     "\n",
-    "    def test_check_balance(self):\n",
-    "        assert_equal(check_balance(None), False)\n",
+    "    @raises(Exception)\n",
+    "    def test_check_balance_empty(self):\n",
+    "        check_balance(None)\n",
     "\n",
+    "    def test_check_balance(self):\n",
     "        node = Node(5)\n",
     "        assert_equal(check_balance(node), True)\n",
     "\n",
@@ -158,6 +161,7 @@
     "\n",
     "def main():\n",
     "    test = TestCheckBalance()\n",
+    "    test.test_check_balance_empty()\n",
     "    test.test_check_balance()\n",
     "\n",
     "\n",
diff --git a/graphs_trees/check_balance/check_balance_solution.ipynb b/graphs_trees/check_balance/check_balance_solution.ipynb
index 775b786..56b7271 100644
--- a/graphs_trees/check_balance/check_balance_solution.ipynb
+++ b/graphs_trees/check_balance/check_balance_solution.ipynb
@@ -35,7 +35,7 @@
     "\n",
     "* Is a balanced tree one where the heights of two sub trees of any node doesn't differ by more than 1?\n",
     "    * Yes\n",
-    "* If this is called on a None input, should we return False?\n",
+    "* If this is called on a None input, should we raise an exception?\n",
     "    * Yes\n",
     "* Can we assume we already have a Node class with an insert method?\n",
     "    * Yes\n",
@@ -104,7 +104,7 @@
    "outputs": [],
    "source": [
     "def _check_balance(root):\n",
-    "    if not root:\n",
+    "    if root is None:\n",
     "        return (True, 0, 0)\n",
     "    left_balanced, left_min_h, left_max_h = _check_balance(root.left)\n",
     "    right_balanced, right_min_h, right_max_h = _check_balance(root.right)\n",
@@ -113,9 +113,10 @@
     "    balanced = left_balanced and right_balanced and abs(max_h-min_h) <= 1\n",
     "    return (balanced, min_h, max_h)\n",
     "\n",
+    "\n",
     "def check_balance(root):\n",
     "    if root is None:\n",
-    "        return False\n",
+    "        raise Exception('No root node')\n",
     "    balanced, _, _ = _check_balance(root)\n",
     "    return balanced"
    ]
@@ -145,13 +146,16 @@
    "source": [
     "%%writefile test_check_balance.py\n",
     "from nose.tools import assert_equal\n",
+    "from nose.tools import raises\n",
     "\n",
     "\n",
     "class TestCheckBalance(object):\n",
     "\n",
-    "    def test_check_balance(self):\n",
-    "        assert_equal(check_balance(None), False)\n",
+    "    @raises(Exception)\n",
+    "    def test_check_balance_empty(self):\n",
+    "        check_balance(None)\n",
     "\n",
+    "    def test_check_balance(self):\n",
     "        node = Node(5)\n",
     "        assert_equal(check_balance(node), True)\n",
     "\n",
@@ -182,6 +186,7 @@
     "\n",
     "def main():\n",
     "    test = TestCheckBalance()\n",
+    "    test.test_check_balance_empty()\n",
     "    test.test_check_balance()\n",
     "\n",
     "\n",
diff --git a/graphs_trees/check_balance/test_check_balance.py b/graphs_trees/check_balance/test_check_balance.py
index 843ba6c..13dafd4 100644
--- a/graphs_trees/check_balance/test_check_balance.py
+++ b/graphs_trees/check_balance/test_check_balance.py
@@ -1,11 +1,14 @@
 from nose.tools import assert_equal
+from nose.tools import raises
 
 
 class TestCheckBalance(object):
 
-    def test_check_balance(self):
-        assert_equal(check_balance(None), False)
+    @raises(Exception)
+    def test_check_balance_empty(self):
+        check_balance(None)
 
+    def test_check_balance(self):
         node = Node(5)
         assert_equal(check_balance(node), True)
 
@@ -36,6 +39,7 @@ class TestCheckBalance(object):
 
 def main():
     test = TestCheckBalance()
+    test.test_check_balance_empty()
     test.test_check_balance()