Cleanup class decorator. Closes #363.

moto/core/models.py

@@ -18,10 +18,10 @@ class MockAWS(object):
         if self.__class__.nested_count == 0:
             HTTPretty.reset()
 
-    def __call__(self, func):
+    def __call__(self, func, reset=True):
         if inspect.isclass(func):
             return self.decorate_class(func)
-        return self.decorate_callable(func)
+        return self.decorate_callable(func, reset)
 
     def __enter__(self):
         self.start()
@@ -64,10 +64,13 @@ class MockAWS(object):
             HTTPretty.disable()
             HTTPretty.reset()
 
-    def decorate_callable(self, func):
+    def decorate_callable(self, func, reset):
         def wrapper(*args, **kwargs):
-            with self:
+            self.start(reset=reset)
+            try:
                 result = func(*args, **kwargs)
+            finally:
+                self.stop()
             return result
         functools.update_wrapper(wrapper, func)
         wrapper.__wrapped__ = func
@@ -87,7 +90,7 @@ class MockAWS(object):
                 continue
 
             try:
-                setattr(klass, attr, self(attr_value))
+                setattr(klass, attr, self(attr_value, reset=False))
             except TypeError:
                 # Sometimes we can't set this for built-in types
                 continue

tests/test_core/test_decorator_calls.py

@@ -2,10 +2,12 @@ from __future__ import unicode_literals
 import boto
 from boto.exception import EC2ResponseError
 import sure  # noqa
+import unittest
+
 import tests.backport_assert_raises  # noqa
 from nose.tools import assert_raises
 
-from moto import mock_ec2
+from moto import mock_ec2, mock_s3
 
 '''
 Test the different ways that the decorator can be used
@@ -69,3 +71,14 @@ class Tester(object):
     def test_still_the_same(self):
         conn = boto.connect_ec2()
         list(conn.get_all_instances()).should.have.length_of(0)
+
+
+@mock_s3
+class TesterWithSetup(unittest.TestCase):
+    def setUp(self):
+        self.conn = boto.connect_s3()
+        self.conn.create_bucket('mybucket')
+
+    def test_still_the_same(self):
+        bucket = self.conn.get_bucket('mybucket')
+        bucket.name.should.equal("mybucket")