diff --git a/moto/core/models.py b/moto/core/models.py index 9fe1e96bd..1b0f2d0dd 100644 --- a/moto/core/models.py +++ b/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__() diff --git a/moto/sqs/models.py b/moto/sqs/models.py index 1404ded75..f2e3ed400 100644 --- a/moto/sqs/models.py +++ b/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) diff --git a/tests/test_core/test_context_manager.py b/tests/test_core/test_context_manager.py new file mode 100644 index 000000000..4824e021f --- /dev/null +++ b/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'])