Merge pull request #41 from Miserlou/decorator

Adds @placebo_session decorator

README.md

@@ -97,6 +97,8 @@ lambda.list_functions()
 ... mocked response will be returned
 ```
 
+#### Manual Mocking
+
 You can also add mocked responses manually:
 
 ```
@@ -123,3 +125,40 @@ pill.save_response(service='lambda', operation='ListFunctions',
 You can add additional responses to a particular operation and the responses
 will be returned in order.  The final parameter is the HTTP response code which
 is optional.  The default value is 200.
+
+#### Usage as a decorator
+
+Placebo also provides a decorator for easier usage.
+
+First, you'll want to decorate your test method with `placebo_session` and include the `session` kwarg in your method, ex:
+```python
+@placebo_session
+def test_your_function(self, session):
+    foo = Foo()
+    arn = foo.create_iam_roles(session)
+    self.assertEqual(arn, "arn:aws:iam::123:role/{}".format(foo.role_name))
+```
+
+Now, you'll be able to record the AWS interactions with an environment variable:
+```bash
+$ PLACEBO_MODE=record nosetests tests.tests:TestFoo.test_create_iam_roles
+```
+
+You can optionally pass an AWS profile to use:
+```bash
+$ PLACEBO_PROFILE=foo PLACEBO_MODE=record nosetests tests.tests:TestFoo.test_create_iam_roles
+```
+
+In this example, it has created the following JSON blobs:
+```bash
+tests/placebo/TestFoo.test_create_iam_roles
+tests/placebo/TestFoo.test_create_iam_roles/iam.CreateRole_1.json
+tests/placebo/TestFoo.test_create_iam_roles/iam.GetRole_1.json
+tests/placebo/TestFoo.test_create_iam_roles/iam.GetRolePolicy_1.json
+tests/placebo/TestFoo.test_create_iam_roles/iam.PutRolePolicy_1.json
+```
+
+After the JSON has been created, simply drop the environment variables and re-run your test:
+```bash
+$ nosetests tests.tests:TestFoo.test_create_iam_roles
+```

placebo/utils.py

@@ -0,0 +1,48 @@
+import placebo
+import boto3
+import os
+import functools
+
+PLACEBO_DIR = os.path.join(os.path.dirname(__file__), 'placebo')
+
+
+def placebo_session(function):
+    """
+    Decorator to help do testing with placebo.
+    Simply wrap the function you want to test and make sure to add
+    a "session" argument so the decorator can pass the placebo session.
+    Accepts the following environment variables to configure placebo:
+    PLACEBO_MODE: set to "record" to record AWS calls and save them
+    PLACEBO_PROFILE: optionally set an AWS credential profile to record with
+    """
+
+    @functools.wraps(function)
+    def wrapper(*args, **kwargs):
+        session_kwargs = {
+            'region_name': os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')
+        }
+        profile_name = os.environ.get('PLACEBO_PROFILE', None)
+        if profile_name:
+            session_kwargs['profile_name'] = profile_name
+
+        session = boto3.Session(**session_kwargs)
+
+        self = args[0]
+        prefix = self.__class__.__name__ + '.' + function.__name__
+        record_dir = os.path.join(PLACEBO_DIR, prefix)
+
+        if not os.path.exists(record_dir):
+            os.makedirs(record_dir)
+
+        pill = placebo.attach(session, data_path=record_dir)
+
+        if os.environ.get('PLACEBO_MODE') == 'record':
+            pill.record()
+        else:
+            pill.playback()
+
+        kwargs['session'] = session
+
+        return function(*args, **kwargs)
+
+    return wrapper

tests/unit/test_utils.py

@@ -0,0 +1,33 @@
+import unittest
+import os
+import shutil
+
+import boto3
+import mock
+
+from placebo.utils import placebo_session
+
+class TestUtils(unittest.TestCase):
+
+    def setUp(self):
+        self.environ = {}
+        self.environ_patch = mock.patch('os.environ', self.environ)
+        self.environ_patch.start()
+        credential_path = os.path.join(os.path.dirname(__file__), 'cfg',
+                                       'aws_credentials')
+        self.environ['AWS_SHARED_CREDENTIALS_FILE'] = credential_path
+        self.environ['PLACEBO_MODE'] = 'record'
+        self.session = boto3.Session(profile_name='foobar',
+                                     region_name='us-west-2')
+
+    @placebo_session
+    def test_decorator(self, session):
+
+        # Tear it up..
+        PLACEBO_TEST_DIR = os.path.join(os.getcwd(),'placebo', 'placebo')
+        prefix = 'TestUtils.test_decorator'
+        record_dir = os.path.join(PLACEBO_TEST_DIR, prefix)
+        self.assertTrue(os.path.exists(record_dir))
+
+        # Tear it down..
+        shutil.rmtree(PLACEBO_TEST_DIR)