Merge pull request #2291 from spulec/context-manager-return-mock

Have context manager return mock.

moto/core/models.py

@@ -52,6 +52,7 @@ class BaseMockAWS(object):
 
     def __enter__(self):
         self.start()
+        return self
 
     def __exit__(self, *args):
         self.stop()
@@ -465,10 +466,14 @@ class BaseModel(object):
 
 class BaseBackend(object):
 
-    def reset(self):
+    def _reset_model_refs(self):
+        # Remove all references to the models stored
         for service, models in model_data.items():
             for model_name, model in models.items():
                 model.instances = []
+
+    def reset(self):
+        self._reset_model_refs()
         self.__dict__ = {}
         self.__init__()
 

moto/sqs/models.py

@@ -379,6 +379,7 @@ class SQSBackend(BaseBackend):
 
     def reset(self):
         region_name = self.region_name
+        self._reset_model_refs()
         self.__dict__ = {}
         self.__init__(region_name)
 

tests/test_core/test_context_manager.py

@@ -0,0 +1,12 @@
+import sure  # noqa
+import boto3
+from moto import mock_sqs, settings
+
+
+def test_context_manager_returns_mock():
+    with mock_sqs() as sqs_mock:
+        conn = boto3.client("sqs", region_name='us-west-1')
+        conn.create_queue(QueueName="queue1")
+
+        if not settings.TEST_SERVER_MODE:
+            list(sqs_mock.backends['us-west-1'].queues.keys()).should.equal(['queue1'])