port tests to python 2.6

1 files changed, 5 insertions, 1 deletions

diff --git a/tests/__init__.py b/tests/__init__.py
index 757e47d..5843e5f 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -1,9 +1,13 @@
 import aiogreen
-import unittest
 try:
     import asyncio
 except ImportError:
     import trollius as asyncio
+try:
+    # On Python 2.6, unittest2 is needed to get new features like addCleanup()
+    import unittest2 as unittest
+except ImportError:
+    import unittest
 
 class TestCase(unittest.TestCase):
     def setUp(self):