Handle unittest2 test suites as well.

This doesn't address the incompat with setUpClass etc, but I have
documented it now.

2 files changed, 26 insertions, 2 deletions

diff --git a/testresources/__init__.py b/testresources/__init__.py
index 2b814f1..f5dbaf0 100644
--- a/testresources/__init__.py
+++ b/testresources/__init__.py
@@ -20,6 +20,10 @@
 import heapq
 import inspect
 import unittest
+try:
+    import unittest2
+except ImportError:
+    unittest2 = None
 
 # same format as sys.version_info: "A tuple containing the five components of
 # the version number: major, minor, micro, releaselevel, and serial. All
@@ -190,6 +194,8 @@ def _strongly_connected_components(graph, no_resources):
 class OptimisingTestSuite(unittest.TestSuite):
     """A resource creation optimising TestSuite."""
 
+    known_suite_classes = None
+
     def adsorbSuite(self, test_case_or_suite):
         """Deprecated. Use addTest instead."""
         self.addTest(test_case_or_suite)
@@ -207,8 +213,7 @@ class OptimisingTestSuite(unittest.TestSuite):
         except TypeError:
             unittest.TestSuite.addTest(self, test_case_or_suite)
             return
-        if test_case_or_suite.__class__ in (unittest.TestSuite,
-                                            OptimisingTestSuite):
+        if test_case_or_suite.__class__ in self.__class__.known_suite_classes:
             for test in tests:
                 self.adsorbSuite(test)
         else:
@@ -413,6 +418,12 @@ class OptimisingTestSuite(unittest.TestSuite):
         return order[1:]
 
 
+OptimisingTestSuite.known_suite_classes = (
+    unittest.TestSuite, OptimisingTestSuite)
+if unittest2 is not None:
+    OptimisingTestSuite.known_suite_classes += (unittest2.TestSuite,)
+
+
 class TestLoader(unittest.TestLoader):
     """Custom TestLoader to set the right TestSuite class."""
     suiteClass = OptimisingTestSuite
diff --git a/testresources/tests/test_optimising_test_suite.py b/testresources/tests/test_optimising_test_suite.py
index c319f1f..8293885 100644
--- a/testresources/tests/test_optimising_test_suite.py
+++ b/testresources/tests/test_optimising_test_suite.py
@@ -21,6 +21,10 @@ import testresources
 from testresources import split_by_resources
 from testresources.tests import ResultWithResourceExtensions
 import unittest
+try:
+    import unittest2
+except ImportError:
+    unittest2 = None
 
 
 def test_suite():
@@ -98,6 +102,15 @@ class TestOptimisingTestSuite(testtools.TestCase):
         self.optimising_suite.addTest(suite)
         self.assertEqual([case], self.optimising_suite._tests)
 
+    @testtools.skipIf(unittest2 is None, "Unittest2 needed")
+    def testAddUnittest2TestSuite(self):
+        # Adding a unittest2 test suite is the same as adding all the tests in
+        # that suite.
+        case = self.makeTestCase()
+        suite = unittest2.TestSuite([case])
+        self.optimising_suite.addTest(suite)
+        self.assertEqual([case], self.optimising_suite._tests)
+
     def testAddTestOptimisingTestSuite(self):
         # when adding an optimising test suite, it should be unpacked.
         case = self.makeTestCase()