Merge adsorbSuite->addTests rename.

2 files changed, 67 insertions, 29 deletions

diff --git a/lib/testresources/__init__.py b/lib/testresources/__init__.py
index 2f317ea..5a5aa15 100644
--- a/lib/testresources/__init__.py
+++ b/lib/testresources/__init__.py
@@ -25,18 +25,6 @@ def test_suite():
     return testresources.tests.test_suite()
 
 
-def iterate_tests(test_suite_or_case):
-    """Iterate through all of the test cases in `test_suite_or_case`."""
-    try:
-        suite = iter(test_suite_or_case)
-    except TypeError:
-        yield test_suite_or_case
-    else:
-        for test in suite:
-            for subtest in iterate_tests(test):
-                yield subtest
-
-
 def split_by_resources(tests):
     """Split a list of tests by the resources that the tests use.
 
@@ -60,14 +48,29 @@ class OptimisingTestSuite(unittest.TestSuite):
     """A resource creation optimising TestSuite."""
 
     def adsorbSuite(self, test_case_or_suite):
-        """Add `test_case_or_suite`, unwrapping any suites we find.
+        """Deprecated. Use addTest instead."""
+        self.addTest(test_case_or_suite)
+
+    def addTest(self, test_case_or_suite):
+        """Add `test_case_or_suite`, unwrapping standard TestSuites.
 
-        This means that any containing TestSuites will be removed. These
-        suites might have their own unittest extensions, so be careful with
-        this.
+        This means that any containing unittest.TestSuites will be removed,
+        while any custom test suites will be 'distributed' across their
+        members. Thus addTest(CustomSuite([a, b])) will result in
+        CustomSuite([a]) and CustomSuite([b]) being added to this suite.
         """
-        for test in iterate_tests(test_case_or_suite):
-            self.addTest(test)
+        try:
+            tests = iter(test_case_or_suite)
+        except TypeError:
+            unittest.TestSuite.addTest(self, test_case_or_suite)
+            return
+        if unittest.TestSuite == test_case_or_suite.__class__:
+            for test in tests:
+                self.adsorbSuite(test)
+        else:
+            for test in tests:
+                unittest.TestSuite.addTest(
+                    self, test_case_or_suite.__class__([test]))
 
     def cost_of_switching(self, old_resource_set, new_resource_set):
         """Cost of switching from 'old_resource_set' to 'new_resource_set'.
diff --git a/lib/testresources/tests/test_optimising_test_suite.py b/lib/testresources/tests/test_optimising_test_suite.py
index 4bf06d3..fddbb5f 100644
--- a/lib/testresources/tests/test_optimising_test_suite.py
+++ b/lib/testresources/tests/test_optimising_test_suite.py
@@ -32,6 +32,16 @@ def test_suite():
     return result
 
 
+class CustomSuite(unittest.TestSuite):
+    """Custom TestSuite that's comparable using == and !=."""
+
+    def __eq__(self, other):
+        return (self.__class__ == other.__class__
+                and self._tests == other._tests)
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
 class MakeCounter(testresources.TestResource):
     """Test resource that counts makes and cleans."""
 
@@ -71,33 +81,58 @@ class TestOptimisingTestSuite(testtools.TestCase):
         testtools.TestCase.setUp(self)
         self.optimising_suite = testresources.OptimisingTestSuite()
 
-    def testAdsorbTest(self):
-        # Adsorbing a single test case is the same as adding one using
-        # addTest.
+    def testAddTest(self):
+        # Adding a single test case is the same as adding one using the
+        # standard addTest.
         case = self.makeTestCase()
-        self.optimising_suite.adsorbSuite(case)
+        self.optimising_suite.addTest(case)
         self.assertEqual([case], self.optimising_suite._tests)
 
-    def testAdsorbTestSuite(self):
-        # Adsorbing a test suite will is the same as adding all the tests in
+    def testAddTestSuite(self):
+        # Adding a standard test suite is the same as adding all the tests in
         # that suite.
         case = self.makeTestCase()
         suite = unittest.TestSuite([case])
-        self.optimising_suite.adsorbSuite(suite)
+        self.optimising_suite.addTest(suite)
         self.assertEqual([case], self.optimising_suite._tests)
 
-    def testAdsorbFlattensAllSuiteStructure(self):
-        # adsorbSuite will get rid of all suite structure when adding a test,
-        # no matter how much nesting is going on.
+    def testAddFlattensStandardSuiteStructure(self):
+        # addTest will get rid of all unittest.TestSuite structure when adding
+        # a test, no matter how much nesting is going on.
         case1 = self.makeTestCase()
         case2 = self.makeTestCase()
         case3 = self.makeTestCase()
         suite = unittest.TestSuite(
             [unittest.TestSuite([case1, unittest.TestSuite([case2])]),
              case3])
-        self.optimising_suite.adsorbSuite(suite)
+        self.optimising_suite.addTest(suite)
         self.assertEqual([case1, case2, case3], self.optimising_suite._tests)
 
+    def testAddDistributesNonStandardSuiteStructure(self):
+        # addTest distributes all non-standard TestSuites across their
+        # members.
+        case1 = self.makeTestCase()
+        case2 = self.makeTestCase()
+        inner_suite = unittest.TestSuite([case2])
+        suite = CustomSuite([case1, inner_suite])
+        self.optimising_suite.addTest(suite)
+        self.assertEqual(
+            [CustomSuite([case1]), CustomSuite([inner_suite])],
+            self.optimising_suite._tests)
+
+    def testAddPullsNonStandardSuitesUp(self):
+        # addTest flattens standard TestSuites, even those that contain custom
+        # suites. When it reaches the custom suites, it distributes them
+        # across their members.
+        case1 = self.makeTestCase()
+        case2 = self.makeTestCase()
+        inner_suite = CustomSuite([case1, case2])
+        self.optimising_suite.addTest(
+            unittest.TestSuite([unittest.TestSuite([inner_suite])]))
+        self.assertEqual(
+            [CustomSuite([case1]), CustomSuite([case2])],
+            self.optimising_suite._tests)
+
     def testSingleCaseResourceAcquisition(self):
         sample_resource = MakeCounter()
         def getResourceCount():