Writing a custom test loader to load tests by attribute

Author: Siddharta Govindaraj

run_test.py

@@ -1,7 +1,24 @@
 import unittest
 
-from stock_alerter.tests import test_stock
+
+class AttribLoader(unittest.TestLoader):
+    def __init__(self, attrib):
+        self.attrib = attrib
+
+    def loadTestsFromModule(self, module, use_load_tests=False):
+        return super().loadTestsFromModule(module, use_load_tests=False)
+
+    def getTestCaseNames(self, testCaseClass):
+        test_names = super().getTestCaseNames(testCaseClass)
+        filtered_test_names = [test
+                               for test in test_names
+                               if hasattr(getattr(testCaseClass, test),
+                                          self.attrib)]
+        return filtered_test_names
+
 
 if __name__ == "__main__":
+    loader = AttribLoader("slow")
+    test_suite = loader.discover(".")
     runner = unittest.TextTestRunner()
-    runner.run(test_stock.suite())
+    runner.run(test_suite)

stock_alerter/tests/test_stock.py

@@ -1,4 +1,3 @@
-import sys
 import unittest
 import collections
 from datetime import datetime, timedelta
@@ -13,7 +12,6 @@ def setUp(self):
     def test_price_of_a_new_stock_class_should_be_None(self):
         self.assertIsNone(self.goog.price)
 
-    @unittest.skipIf(sys.platform.startswith("win"), "skip on windows")
     def test_stock_update(self):
         """An update should set the price on the stock object
 
@@ -22,7 +20,8 @@ def test_stock_update(self):
         self.goog.update(datetime(2014, 2, 12), price=10)
         self.assertEqual(10, self.goog.price)
 
-    @unittest.skipUnless(sys.platform.startswith("win"), "only run on windows")
+    test_stock_update.slow = True
+
     def test_negative_price_should_throw_ValueError(self):
         with self.assertRaises(ValueError):
             self.goog.update(datetime(2014, 2, 13), -1)