Support countTestCases in TestWithScenarios.

2 files changed, 40 insertions, 0 deletions

diff --git a/lib/testscenarios/testcase.py b/lib/testscenarios/testcase.py
index 712d543..23907fb 100644
--- a/lib/testscenarios/testcase.py
+++ b/lib/testscenarios/testcase.py
@@ -37,6 +37,13 @@ class TestWithScenarios(unittest.TestCase):
     subclass (or overridden compatibly with TestWithScenarios).
     """
 
+    def countTestCases(self):
+        scenarios = getattr(self, 'scenarios', None)
+        if not scenarios:
+            return 1
+        else:
+            return len(scenarios)
+
     def run(self, result=None):
         scenarios = getattr(self, 'scenarios', None)
         if scenarios:
diff --git a/lib/testscenarios/tests/test_testcase.py b/lib/testscenarios/tests/test_testcase.py
index 5b86e24..8453c18 100644
--- a/lib/testscenarios/tests/test_testcase.py
+++ b/lib/testscenarios/tests/test_testcase.py
@@ -98,3 +98,36 @@ class TestTestWithScenarios(unittest.TestCase):
         self.assertNotEqual(None, test.scenarios)
         self.assertEqual(None, log[0][1].scenarios)
         self.assertEqual(None, log[4][1].scenarios)
+
+    def test_countTestCases_no_scenarios(self):
+        class ReferenceTest(testscenarios.TestWithScenarios):
+            def test_check_foo(self):
+                pass
+        test = ReferenceTest("test_check_foo")
+        self.assertEqual(1, test.countTestCases())
+
+    def test_countTestCases_empty_scenarios(self):
+        class ReferenceTest(testscenarios.TestWithScenarios):
+            scenarios = []
+            def test_check_foo(self):
+                pass
+        test = ReferenceTest("test_check_foo")
+        self.assertEqual(1, test.countTestCases())
+
+    def test_countTestCases_1_scenarios(self):
+        class ReferenceTest(testscenarios.TestWithScenarios):
+            scenarios = [('1', {'foo': 1, 'bar': 2})]
+            def test_check_foo(self):
+                pass
+        test = ReferenceTest("test_check_foo")
+        self.assertEqual(1, test.countTestCases())
+
+    def test_countTestCases_2_scenarios(self):
+        class ReferenceTest(testscenarios.TestWithScenarios):
+            scenarios = [
+                ('1', {'foo': 1, 'bar': 2}),
+                ('2', {'foo': 2, 'bar': 4})]
+            def test_check_foo(self):
+                pass
+        test = ReferenceTest("test_check_foo")
+        self.assertEqual(2, test.countTestCases())