From db943bcdbb420d07a673cfb6892a4e44b93b388c Mon Sep 17 00:00:00 2001
From: Steve Pulec <spulec@gmail.com>
Date: Wed, 27 Feb 2013 22:25:15 -0500
Subject: [PATCH] add other ways to call decorator

---
 README.md                               | 47 ++++++++++++++++++-
 moto/core/models.py                     | 61 +++++++++++++++++--------
 tests/test_core/test_decorator_calls.py | 39 ++++++++++++++++
 3 files changed, 127 insertions(+), 20 deletions(-)
 create mode 100644 tests/test_core/test_decorator_calls.py

diff --git a/README.md b/README.md
index 4b5447e8..483fad2a 100644
--- a/README.md
+++ b/README.md
@@ -55,4 +55,49 @@ With the decorator wrapping the test, all the calls to s3 are automatically mock
 ```console
 pip install moto
 ```
- -->
\ No newline at end of file
+ -->
+
+## 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()
+```
diff --git a/moto/core/models.py b/moto/core/models.py
index 48a695c6..6c3c8f50 100644
--- a/moto/core/models.py
+++ b/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)
diff --git a/tests/test_core/test_decorator_calls.py b/tests/test_core/test_decorator_calls.py
new file mode 100644
index 00000000..4b58916a
--- /dev/null
+++ b/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)