From 3ed9428cb01cc0af157dfaf50fc4550e1a6d3ab5 Mon Sep 17 00:00:00 2001
From: Steve Pulec <spulec@gmail.com>
Date: Sat, 27 Jun 2015 19:01:01 -0400
Subject: [PATCH] Cleanup class decorator. Closes #363.

---
 moto/core/models.py                     | 13 ++++++++-----
 tests/test_core/test_decorator_calls.py | 15 ++++++++++++++-
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/moto/core/models.py b/moto/core/models.py
index be25321c7..60e744ff5 100644
--- a/moto/core/models.py
+++ b/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
diff --git a/tests/test_core/test_decorator_calls.py b/tests/test_core/test_decorator_calls.py
index bfe5d5730..7d32bc8b3 100644
--- a/tests/test_core/test_decorator_calls.py
+++ b/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")