diff --git a/arrays-strings/compress.ipynb b/arrays-strings/compress.ipynb
index 6c3eea1..2e58357 100644
--- a/arrays-strings/compress.ipynb
+++ b/arrays-strings/compress.ipynb
@@ -49,6 +49,28 @@
     "* 'AAABCCDDDD' -> 'A3B1C2D4'"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "from nose.tools import assert_equal\n",
+    "\n",
+    "class Test(object):\n",
+    "    def test_compress(self, func):\n",
+    "        assert_equal(func(None), None)\n",
+    "        assert_equal(func(''), '')\n",
+    "        assert_equal(func('ABC'), 'ABC')\n",
+    "        assert_equal(func('AAABCCDDDD'), 'A3B1C2D4')\n",
+    "\n",
+    "def run_tests(func):\n",
+    "    test = Test()\n",
+    "    test.test_compress(func)"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -136,14 +158,7 @@
     "    compressed_string.append(str(count))\n",
     "    return \"\".join(compressed_string)\n",
     "\n",
-    "string0 = None\n",
-    "string1 = ''\n",
-    "string2 = 'ABC'\n",
-    "string3 = 'AAABCCDDDD'\n",
-    "print(compress_string(string0))\n",
-    "print(compress_string(string1))\n",
-    "print(compress_string(string2))\n",
-    "print(compress_string(string3))"
+    "run_tests(compress_string)"
    ]
   },
   {
@@ -178,7 +193,7 @@
    },
    "outputs": [],
    "source": [
-    "def compress_string(string):\n",
+    "def compress_string_alt(string):\n",
     "    if string is None or len(string) == 0:\n",
     "        return string\n",
     "    size = 0\n",
@@ -211,14 +226,7 @@
     "    compressed_string[pos + 1] = ord(str(count))\n",
     "    return compressed_string\n",
     "\n",
-    "string0 = None\n",
-    "string1 = bytearray('')\n",
-    "string2 = bytearray('ABC')\n",
-    "string3 = bytearray('AAABCCDDDD')\n",
-    "print(compress_string(string0))\n",
-    "print(compress_string(string1))\n",
-    "print(compress_string(string2))\n",
-    "print(compress_string(string3))"
+    "run_tests(compress_string_alt)"
    ]
   }
  ],