add other ways to call decorator

moto/core/models.py

@@ -4,6 +4,43 @@ import re
 from moto.packages.httpretty import HTTPretty
 
 
+class MockAWS(object):
+    def __init__(self, backend):
+        self.backend = backend
+
+    def __call__(self, func):
+        return self.decorate_callable(func)
+
+    def __enter__(self):
+        self.start()
+
+    def __exit__(self, *args):
+        self.stop()
+
+    def start(self):
+        self.backend.reset()
+        HTTPretty.enable()
+
+        for method in HTTPretty.METHODS:
+            for key, value in self.backend.urls.iteritems():
+                HTTPretty.register_uri(
+                    method=method,
+                    uri=re.compile(key),
+                    body=value,
+                )
+
+    def stop(self):
+        HTTPretty.disable()
+
+    def decorate_callable(self, func):
+        def wrapper(*args, **kwargs):
+            with self:
+                result = func(*args, **kwargs)
+            return result
+        functools.update_wrapper(wrapper, func)
+        return wrapper
+
+
 class BaseBackend(object):
 
     def reset(self):
@@ -18,22 +55,8 @@ class BaseBackend(object):
         urls = backend_urls_module.urls
         return urls
 
-    def decorator(self, func):
-        @functools.wraps(func)
-        def wrapper(*args, **kw):
-            self.reset()
-
-            HTTPretty.enable()
-
-            for method in HTTPretty.METHODS:
-                for key, value in self.urls.iteritems():
-                    HTTPretty.register_uri(
-                        method=method,
-                        uri=re.compile(key),
-                        body=value,
-                    )
-            try:
-                return func(*args, **kw)
-            finally:
-                HTTPretty.disable()
-        return wrapper
+    def decorator(self, func=None):
+        if func:
+            return MockAWS(self)(func)
+        else:
+            return MockAWS(self)