add other ways to call decorator

README.md

@@ -55,4 +55,49 @@ With the decorator wrapping the test, all the calls to s3 are automatically mock
 ```console
 pip install moto
 ```
- -->
+ -->
+
+## Usage
+
+All of the services can be used as a decorator, context manager, or in a raw form.
+
+### Decorator
+
+```python
+@mock_s3
+def test_my_model_save():
+    model_instance = MyModel('steve', 'is awesome')
+    model_instance.save()
+
+    conn = boto.connect_s3()
+    assert conn.get_bucket('mybucket').get_key('steve') == 'is awesome'
+```
+
+### Context Manager
+
+```python
+def test_my_model_save():
+    with mock_s3():
+        model_instance = MyModel('steve', 'is awesome')
+        model_instance.save()
+
+        conn = boto.connect_s3()
+        assert conn.get_bucket('mybucket').get_key('steve') == 'is awesome'
+```
+
+
+### Raw use
+
+```python
+def test_my_model_save():
+    mock = mock_s3()
+    mock.start()
+
+    model_instance = MyModel('steve', 'is awesome')
+    model_instance.save()
+
+    conn = boto.connect_s3()
+    assert conn.get_bucket('mybucket').get_key('steve') == 'is awesome'
+
+    mock.stop()
+```

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)

tests/test_core/test_decorator_calls.py

@@ -0,0 +1,39 @@
+import boto
+from boto.exception import EC2ResponseError
+from sure import expect
+
+from moto import mock_ec2
+
+'''
+Test the different ways that the decorator can be used
+'''
+
+
+@mock_ec2
+def test_basic_decorator():
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    list(conn.get_all_instances()).should.equal([])
+
+
+def test_context_manager():
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    conn.get_all_instances.when.called_with().should.throw(EC2ResponseError)
+
+    with mock_ec2():
+        conn = boto.connect_ec2('the_key', 'the_secret')
+        list(conn.get_all_instances()).should.equal([])
+
+    conn.get_all_instances.when.called_with().should.throw(EC2ResponseError)
+
+
+def test_decorator_start_and_stop():
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    conn.get_all_instances.when.called_with().should.throw(EC2ResponseError)
+
+    mock = mock_ec2()
+    mock.start()
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    list(conn.get_all_instances()).should.equal([])
+    mock.stop()
+
+    expect(conn.get_all_instances.when.called_with()).should.throw(EC2ResponseError)