From b065a20d88daaae0cd6ef3a90bcf7210d97ddf51 Mon Sep 17 00:00:00 2001 From: Bert Blommers Date: Mon, 22 Nov 2021 22:06:59 -0100 Subject: [PATCH] Rework mock_all to work as a context manager (#4607) --- moto/__init__.py | 34 +++++++++++--------- moto/xray/__init__.py | 3 +- moto/xray/mock_client.py | 48 ++++++++++++++++++----------- tests/test_core/test_mock_all.py | 19 ++++++++++-- tests/test_xray/test_xray_client.py | 35 +++++++++++++++++++++ 5 files changed, 104 insertions(+), 35 deletions(-) diff --git a/moto/__init__.py b/moto/__init__.py index 356ab65cf..b4e7527a0 100644 --- a/moto/__init__.py +++ b/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 diff --git a/moto/xray/__init__.py b/moto/xray/__init__.py index 381d9da51..e47d7642b 100644 --- a/moto/xray/__init__.py +++ b/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() diff --git a/moto/xray/mock_client.py b/moto/xray/mock_client.py index 9e042c594..d7e06279f 100644 --- a/moto/xray/mock_client.py +++ b/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): diff --git a/tests/test_core/test_mock_all.py b/tests/test_core/test_mock_all.py index bed54c72b..1a7e07f83 100644 --- a/tests/test_core/test_mock_all.py +++ b/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() diff --git a/tests/test_xray/test_xray_client.py b/tests/test_xray/test_xray_client.py index f28fcc064..81ab70f05 100644 --- a/tests/test_xray/test_xray_client.py +++ b/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