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')