From fb0223bfc5c1a489a62124927d7a129be8be67fb Mon Sep 17 00:00:00 2001
From: Chris Martin <ch.martin@gmail.com>
Date: Fri, 23 May 2014 14:24:15 -0400
Subject: [PATCH] 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                  | 46 +++++++++++++++++++-
 tests/functional/test_decorator.py | 69 +++++++++++++++++++++++++++++-
 2 files changed, 112 insertions(+), 3 deletions(-)

diff --git a/httpretty/core.py b/httpretty/core.py
index 5e7e96b..282718a 100644
--- a/httpretty/core.py
+++ b/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):
diff --git a/tests/functional/test_decorator.py b/tests/functional/test_decorator.py
index 65f6f2a..dacb03c 100644
--- a/tests/functional/test_decorator.py
+++ b/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')
\ No newline at end of file
+        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')