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)