Rework mock_all to work as a context manager (#4607)

moto/__init__.py

@@ -1,5 +1,6 @@
 import importlib
 import sys
+from contextlib import ContextDecorator
 
 
 def lazy_load(
@@ -171,22 +172,27 @@ mock_wafv2 = lazy_load(".wafv2", "mock_wafv2")
 mock_sdb = lazy_load(".sdb", "mock_sdb", boto3_name="sdb")
 
 
-def mock_all():
-    dec_names = [
-        d
-        for d in dir(sys.modules["moto"])
-        if d.startswith("mock_")
-        and not d.endswith("_deprecated")
-        and not d == "mock_all"
-    ]
+class MockAll(ContextDecorator):
+    def __init__(self):
+        self.mocks = []
+        for mock in dir(sys.modules["moto"]):
+            if (
+                mock.startswith("mock_")
+                and not mock.endswith("_deprecated")
+                and not mock == ("mock_all")
+            ):
+                self.mocks.append(globals()[mock]())
 
-    def deco(f):
-        for dec_name in reversed(dec_names):
-            dec = globals()[dec_name]
-            f = dec(f)
-        return f
+    def __enter__(self):
+        for mock in self.mocks:
+            mock.start()
 
-    return deco
+    def __exit__(self, *exc):
+        for mock in self.mocks:
+            mock.stop()
+
+
+mock_all = MockAll
 
 
 # import logging

moto/xray/__init__.py

@@ -1,6 +1,7 @@
 from .models import xray_backends
 from ..core.models import base_decorator
-from .mock_client import mock_xray_client, XRaySegment  # noqa
+from .mock_client import MockXrayClient, XRaySegment  # noqa
 
 xray_backend = xray_backends["us-east-1"]
 mock_xray = base_decorator(xray_backends)
+mock_xray_client = MockXrayClient()

moto/xray/mock_client.py

@@ -1,4 +1,3 @@
-from functools import wraps
 import os
 from moto.xray import xray_backends
 import aws_xray_sdk.core
@@ -32,7 +31,7 @@ class MockEmitter(UDPEmitter):
         raise RuntimeError("Should not be running this")
 
 
-def mock_xray_client(f):
+class MockXrayClient:
     """
     Mocks the X-Ray sdk by pwning its evil singleton with our methods
 
@@ -43,30 +42,43 @@ def mock_xray_client(f):
     that itno the recorder instance.
     """
 
-    @wraps(f)
-    def _wrapped(*args, **kwargs):
-        print("Starting X-Ray Patch")
+    def __call__(self, f=None):
+        if not f:
+            return self
 
-        old_xray_context_var = os.environ.get("AWS_XRAY_CONTEXT_MISSING")
+        def wrapped_f():
+            self.start()
+            try:
+                f()
+            finally:
+                self.stop()
+
+        return wrapped_f
+
+    def start(self):
+        print("Starting X-Ray Patch")
+        self.old_xray_context_var = os.environ.get("AWS_XRAY_CONTEXT_MISSING")
         os.environ["AWS_XRAY_CONTEXT_MISSING"] = "LOG_ERROR"
-        old_xray_context = aws_xray_sdk.core.xray_recorder._context
-        old_xray_emitter = aws_xray_sdk.core.xray_recorder._emitter
+        self.old_xray_context = aws_xray_sdk.core.xray_recorder._context
+        self.old_xray_emitter = aws_xray_sdk.core.xray_recorder._emitter
         aws_xray_sdk.core.xray_recorder._context = AWSContext()
         aws_xray_sdk.core.xray_recorder._emitter = MockEmitter()
 
-        try:
-            return f(*args, **kwargs)
-        finally:
+    def stop(self):
+        if self.old_xray_context_var is None:
+            del os.environ["AWS_XRAY_CONTEXT_MISSING"]
+        else:
+            os.environ["AWS_XRAY_CONTEXT_MISSING"] = self.old_xray_context_var
 
-            if old_xray_context_var is None:
-                del os.environ["AWS_XRAY_CONTEXT_MISSING"]
-            else:
-                os.environ["AWS_XRAY_CONTEXT_MISSING"] = old_xray_context_var
+        aws_xray_sdk.core.xray_recorder._emitter = self.old_xray_emitter
+        aws_xray_sdk.core.xray_recorder._context = self.old_xray_context
 
-            aws_xray_sdk.core.xray_recorder._emitter = old_xray_emitter
-            aws_xray_sdk.core.xray_recorder._context = old_xray_context
+    def __enter__(self):
+        self.start()
+        return self
 
-    return _wrapped
+    def __exit__(self, *args):
+        self.stop()
 
 
 class XRaySegment(object):

tests/test_core/test_mock_all.py

@@ -1,14 +1,15 @@
 import boto3
+import pytest
 import sure  # noqa # pylint: disable=unused-import
 
 from moto import mock_all
 
 
 @mock_all()
-def test_multiple_services():
+def test_decorator():
     rgn = "us-east-1"
     sqs = boto3.client("sqs", region_name=rgn)
-    r = sqs.list_queues()  #
+    r = sqs.list_queues()
     r["ResponseMetadata"]["HTTPStatusCode"].should.equal(200)
 
     lmbda = boto3.client("lambda", region_name=rgn)
@@ -18,3 +19,17 @@ def test_multiple_services():
     ddb = boto3.client("dynamodb", region_name=rgn)
     r = ddb.list_tables()
     r["ResponseMetadata"]["HTTPStatusCode"].should.equal(200)
+
+
+def test_context_manager():
+    rgn = "us-east-1"
+
+    with mock_all():
+        sqs = boto3.client("sqs", region_name=rgn)
+        r = sqs.list_queues()
+        r["ResponseMetadata"]["HTTPStatusCode"].should.equal(200)
+
+    unpatched_sqs = boto3.Session().client("sqs", region_name=rgn)
+
+    with pytest.raises(Exception):
+        unpatched_sqs.list_queues()

tests/test_xray/test_xray_client.py

@@ -45,6 +45,41 @@ def test_xray_dynamo_request_id():
     setattr(requests.Session, "prepare_request", original_session_prep_request)
 
 
+def test_xray_dynamo_request_id_with_context_mgr():
+    with mock_xray_client():
+        assert isinstance(xray_core.xray_recorder._emitter, MockEmitter)
+        with mock_dynamodb2():
+            # Could be ran in any order, so we need to tell sdk that its been unpatched
+            xray_core_patcher._PATCHED_MODULES = set()
+            xray_core.patch_all()
+
+            client = boto3.client("dynamodb", region_name="us-east-1")
+
+            with XRaySegment():
+                resp = client.list_tables()
+                resp["ResponseMetadata"].should.contain("RequestId")
+                id1 = resp["ResponseMetadata"]["RequestId"]
+
+            with XRaySegment():
+                client.list_tables()
+                resp = client.list_tables()
+                id2 = resp["ResponseMetadata"]["RequestId"]
+
+            id1.should_not.equal(id2)
+
+            setattr(
+                botocore.client.BaseClient, "_make_api_call", original_make_api_call
+            )
+            setattr(
+                botocore.endpoint.Endpoint, "_encode_headers", original_encode_headers
+            )
+            setattr(requests.Session, "request", original_session_request)
+            setattr(requests.Session, "prepare_request", original_session_prep_request)
+
+    # Verify we have unmocked the xray recorder
+    assert not isinstance(xray_core.xray_recorder._emitter, MockEmitter)
+
+
 @mock_xray_client
 def test_xray_udp_emitter_patched():
     # Could be ran in any order, so we need to tell sdk that its been unpatched