Decorate unittest.TestCase setUp/tearDown methods

When the class being decorated by httprettified inherits from
unittest.TestCase, enable HTTPretty in setUp rather than decorating each
test separately. This lets users register their HTTP mock entries in
setUp if they want.

httpretty/core.py

@@ -991,8 +991,37 @@ class httpretty(HttpBaseClass):
 
 
 def httprettified(test):
-    "A decorator tests that use HTTPretty"
-    def decorate_class(klass):
+
+    def decorate_unittest_TestCase_setUp(klass):
+
+        # Prefer addCleanup (added in python 2.7), but fall back
+        # to using tearDown if it isn't available
+        use_addCleanup = hasattr(klass, 'addCleanup')
+
+        original_setUp = (klass.setUp
+                          if hasattr(klass, 'setUp')
+                          else None)
+        def new_setUp(self):
+            httpretty.enable()
+            if use_addCleanup:
+                self.addCleanup(httpretty.disable)
+            if original_setUp:
+                original_setUp(self)
+        klass.setUp = new_setUp
+
+        if not use_addCleanup:
+            original_tearDown = (klass.setUp
+                                 if hasattr(klass, 'tearDown')
+                                 else None)
+            def new_tearDown(self):
+                httpretty.disable()
+                if original_tearDown:
+                    original_tearDown(self)
+            klass.tearDown = new_tearDown
+
+        return klass
+
+    def decorate_test_methods(klass):
         for attr in dir(klass):
             if not attr.startswith('test_'):
                 continue
@@ -1004,6 +1033,19 @@ def httprettified(test):
             setattr(klass, attr, decorate_callable(attr_value))
         return klass
 
+    def is_unittest_TestCase(klass):
+        try:
+            import unittest
+            return issubclass(klass, unittest.TestCase)
+        except ImportError:
+            return False
+
+    "A decorator for tests that use HTTPretty"
+    def decorate_class(klass):
+        if is_unittest_TestCase(klass):
+            return decorate_unittest_TestCase_setUp(klass)
+        return decorate_test_methods(klass)
+
     def decorate_callable(test):
         @functools.wraps(test)
         def wrapper(*args, **kw):

tests/functional/test_decorator.py

@@ -22,6 +22,34 @@ def test_decor():
     expect(got1).to.equal(b'glub glub')
 
 
+@httprettified
+class DecoratedNonUnitTest(object):
+
+    def test_fail(self):
+        raise AssertionError('Tests in this class should not '
+                             'be executed by the test runner.')
+
+    def test_decorated(self):
+        HTTPretty.register_uri(
+            HTTPretty.GET, "http://localhost/",
+            body="glub glub")
+
+        fd = urllib2.urlopen('http://localhost/')
+        got1 = fd.read()
+        fd.close()
+
+        expect(got1).to.equal(b'glub glub')
+
+
+class NonUnitTestTest(TestCase):
+    """
+    Checks that the test methods in DecoratedNonUnitTest were decorated.
+    """
+
+    def test_decorated(self):
+        DecoratedNonUnitTest().test_decorated()
+
+
 @httprettified
 class ClassDecorator(TestCase):
 
@@ -45,4 +73,43 @@ class ClassDecorator(TestCase):
         got1 = fd.read()
         fd.close()
 
-        expect(got1).to.equal(b'buble buble')
+        expect(got1).to.equal(b'buble buble')
+
+@httprettified
+class ClassDecoratorWithSetUp(TestCase):
+
+    def setUp(self):
+        HTTPretty.register_uri(
+            HTTPretty.GET, "http://localhost/",
+            responses=[
+                HTTPretty.Response("glub glub"),
+                HTTPretty.Response("buble buble"),
+            ])
+
+    def test_decorated(self):
+
+        fd = urllib2.urlopen('http://localhost/')
+        got1 = fd.read()
+        fd.close()
+
+        expect(got1).to.equal(b'glub glub')
+
+        fd = urllib2.urlopen('http://localhost/')
+        got2 = fd.read()
+        fd.close()
+
+        expect(got2).to.equal(b'buble buble')
+
+    def test_decorated2(self):
+
+        fd = urllib2.urlopen('http://localhost/')
+        got1 = fd.read()
+        fd.close()
+
+        expect(got1).to.equal(b'glub glub')
+
+        fd = urllib2.urlopen('http://localhost/')
+        got2 = fd.read()
+        fd.close()
+
+        expect(got2).to.equal(b'buble buble')