Repository: boto/botocore
Pull request #915 (2016-05-10 17:42:54)

Implement context mangager functions for Stubber
It can be useful to use a stubber as a context manager to make code a
bit more readable and to ensure you aren't leaking events.

I was responding to a question on StackOverflow about stubbing responses
in which they were attempting to use mock. While they didn't complain
about having to use `activate`, I realized my sample code didn't look
as clean as their (unsuccessful) attempt at mocking. Since it is such
a tiny change, I figured it would be useful if I just knocked it out.

## Example
```python
import datetime
import botocore.session
from botocore.stub import Stubber


s3 = botocore.session.get_session().create_client('s3')

response = {
    "Owner": {
        "ID": "foo",
        "DisplayName": "bar"
    },
    "Buckets": [{
        "CreationDate": datetime.datetime(2016, 1, 20, 22, 9),
        "Name": "baz"
    }]
}


with Stubber(s3) as stubber:
    stubber.add_response('list_buckets', response, {})
    service_response = s3.list_buckets()

assert service_response == response
```

cc @kyleknap @jamesls 


Patch:
diff --git a/botocore/stub.py b/botocore/stub.py
index d3c4fac1d..255eac635 100644
--- a/botocore/stub.py
+++ b/botocore/stub.py
@@ -63,6 +63,37 @@ class Stubber(object):
 
         service_response = s3.list_objects(Bucket='test-bucket')
         assert service_response == response
+
+
+    This class can also be called as a context manager, which will handle
+    activation / deactivation for you.
+
+    **Example:**
+    ::
+        import datetime
+        import botocore.session
+        from botocore.stub import Stubber
+
+
+        s3 = botocore.session.get_session().create_client('s3')
+
+        response = {
+            "Owner": {
+                "ID": "foo",
+                "DisplayName": "bar"
+            },
+            "Buckets": [{
+                "CreationDate": datetime.datetime(2016, 1, 20, 22, 9),
+                "Name": "baz"
+            }]
+        }
+
+
+        with Stubber(s3) as stubber:
+            stubber.add_response('list_buckets', response, {})
+            service_response = s3.list_buckets()
+
+        assert service_response == response
     """
     def __init__(self, client):
         """
@@ -73,6 +104,13 @@ class Stubber(object):
         self._expected_params_event_id = 'boto_stubber_expected_params'
         self._queue = deque()
 
+    def __enter__(self):
+        self.activate()
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        self.deactivate()
+
     def activate(self):
         """
         Activates the stubber on the client


Test patch:
diff --git a/tests/functional/test_stub.py b/tests/functional/test_stub.py
index 640193267..b40f2dbaa 100644
--- a/tests/functional/test_stub.py
+++ b/tests/functional/test_stub.py
@@ -38,6 +38,14 @@ class TestStubber(unittest.TestCase):
         response = self.client.list_objects(Bucket='foo')
         self.assertEqual(response, service_response)
 
+    def test_context_manager_returns_response(self):
+        service_response = {'ResponseMetadata': {'foo': 'bar'}}
+        self.stubber.add_response('list_objects', service_response)
+
+        with self.stubber:
+            response = self.client.list_objects(Bucket='foo')
+        self.assertEqual(response, service_response)
+
     def test_activated_stubber_errors_with_no_registered_stubs(self):
         self.stubber.activate()
         with self.assertRaises(StubResponseError):
diff --git a/tests/unit/test_stub.py b/tests/unit/test_stub.py
index b57c87022..9bedb9c21 100644
--- a/tests/unit/test_stub.py
+++ b/tests/unit/test_stub.py
@@ -68,6 +68,23 @@ class TestStubber(unittest.TestCase):
         self.event_emitter.unregister.assert_any_call(
             'before-call.*.*', mock.ANY, unique_id=mock.ANY)
 
+    def test_context_manager(self):
+        self.event_emitter = mock.Mock()
+        self.client.meta.events = self.event_emitter
+
+        with self.stubber:
+            # Ensure events are registered in context
+            self.event_emitter.register_first.assert_called_with(
+                'before-parameter-build.*.*', mock.ANY, unique_id=mock.ANY)
+            self.event_emitter.register.assert_called_with(
+                'before-call.*.*', mock.ANY, unique_id=mock.ANY)
+
+        # Ensure events are no longer registered once we leave the context
+        self.event_emitter.unregister.assert_any_call(
+            'before-parameter-build.*.*', mock.ANY, unique_id=mock.ANY)
+        self.event_emitter.unregister.assert_any_call(
+            'before-call.*.*', mock.ANY, unique_id=mock.ANY)
+
     def test_add_response(self):
         response = {'foo': 'bar'}
         self.stubber.add_response('foo', response)