Core: Reduce slowdown mock_all (#6451)

docs/docs/configuration/environment_variables.rst

@@ -24,6 +24,8 @@ The following is a non-exhaustive list of the environment variables that can be
 +-------------------------------+----------+-----------+-------------------------------------------------------------------------------------------------+
 | MOTO_ALLOW_NONEXISTENT_REGION | bool     | False     |                                                                                                 |
 +-------------------------------+----------+-----------+-------------------------------------------------------------------------------------------------+
+| MOTO_ENABLE_ISO_REGIONS       | bool     | False     | Enable the mocking of ISO regions (ISO, ISO-B, ISO-E, ISO-F)                                    |
++-------------------------------+----------+-----------+-------------------------------------------------------------------------------------------------+
 |                               |          |           |                                                                                                 |
 +-------------------------------+----------+-----------+-------------------------------------------------------------------------------------------------+
 | MOTO_S3_CUSTOM_ENDPOINTS      | str      |           | See :ref:`s3`.                                                                                  |

moto/core/base_backend.py

@@ -5,7 +5,7 @@ from functools import lru_cache
 from threading import RLock
 from typing import Any, List, Dict, Optional, ClassVar, TypeVar, Iterator
 from uuid import uuid4
-from moto.settings import allow_unknown_region
+from moto.settings import allow_unknown_region, enable_iso_regions
 from .model_instances import model_data
 from .utils import convert_regex_to_flask_path
 
@@ -56,13 +56,16 @@ class BaseBackend:
         for url_base in url_bases:
             # The default URL_base will look like: http://service.[..].amazonaws.com/...
             # This extension ensures support for the China & ISO regions
-            alt_dns_suffixes = {
-                "cn": "amazonaws.com.cn",
-                "iso": "c2s.ic.gov",
-                "isob": "sc2s.sgov.gov",
-                "isoe": "cloud.adc-e.uk",
-                "isof": "csp.hci.ic.gov",
-            }
+            alt_dns_suffixes = {"cn": "amazonaws.com.cn"}
+            if enable_iso_regions():
+                alt_dns_suffixes.update(
+                    {
+                        "iso": "c2s.ic.gov",
+                        "isob": "sc2s.sgov.gov",
+                        "isoe": "cloud.adc-e.uk",
+                        "isof": "csp.hci.ic.gov",
+                    }
+                )
 
             for url_path, handler in unformatted_paths.items():
                 url = url_path.format(url_base)

moto/core/custom_responses_mock.py

@@ -23,6 +23,11 @@ class CallbackResponse(responses.CallbackResponse):
     Need to subclass so we can change a couple things
     """
 
+    def __eq__(self, other: Any) -> bool:
+        if isinstance(other, CallbackResponse):
+            return self.method == other.method and self.url.pattern == other.url.pattern  # type: ignore
+        return super().__eq__(other)
+
     def get_response(self, request: Any) -> responses.HTTPResponse:
         """
         Need to override this so we can pass decode_content=False

moto/core/responses_custom_registry.py

@@ -1,6 +1,7 @@
 # This will only exist in responses >= 0.17
 import responses
-from typing import Any, List, Tuple, Optional
+from collections import defaultdict
+from typing import Any, Dict, List, Tuple, Optional
 from .custom_responses_mock import CallbackResponse, not_implemented_callback
 
 
@@ -9,15 +10,25 @@ class CustomRegistry(responses.registries.FirstMatchRegistry):
     Custom Registry that returns requests in an order that makes sense for Moto:
      - Implemented callbacks take precedence over non-implemented-callbacks
      - CallbackResponses are not discarded after first use - users can mock the same URL as often as they like
+     - CallbackResponses are persisted in a dictionary, with the request-method as key
+       This reduces the number of possible responses that we need to search
     """
 
+    def __init__(self) -> None:
+        self._registered: Dict[str, List[responses.BaseResponse]] = defaultdict(list)
+
     def add(self, response: responses.BaseResponse) -> responses.BaseResponse:
-        if response not in self.registered:
-            super().add(response)
+        if response not in self._registered[response.method]:
+            self._registered[response.method].append(response)
         return response
 
+    def reset(self) -> None:
+        self._registered.clear()
+
     def find(self, request: Any) -> Tuple[Optional[responses.BaseResponse], List[str]]:
-        all_possibles = responses._default_mock._registry.registered + self.registered
+        all_possibles = responses._default_mock._registry.registered
+        # We don't have to search through all possible methods - only the ones registered for this particular method
+        all_possibles.extend(self._registered[request.method])
         found = []
         match_failed_reasons = []
         for response in all_possibles:

moto/settings.py

@@ -145,3 +145,7 @@ def get_docker_host() -> str:
 
 def get_cognito_idp_user_pool_id_strategy() -> Optional[str]:
     return os.environ.get("MOTO_COGNITO_IDP_USER_POOL_ID_STRATEGY")
+
+
+def enable_iso_regions() -> bool:
+    return os.environ.get("MOTO_ENABLE_ISO_REGIONS", "false").lower() == "true"

tests/test_awslambda/test_lambda.py

@@ -1,7 +1,7 @@
 import base64
 import json
 import os
-from unittest import SkipTest
+from unittest import SkipTest, mock
 import boto3
 import hashlib
 import pytest
@@ -25,6 +25,7 @@ _lambda_region = "us-west-2"
 boto3.setup_default_session(region_name=_lambda_region)
 
 
+@mock.patch.dict("os.environ", {"MOTO_ENABLE_ISO_REGIONS": "true"})
 @pytest.mark.parametrize("region", ["us-west-2", "cn-northwest-1", "us-isob-east-1"])
 @mock_lambda
 def test_lambda_regions(region):

tests/test_stepfunctions/test_stepfunctions.py

@@ -888,6 +888,7 @@ def test_state_machine_get_execution_history_contains_expected_success_events_wh
     execution_history["events"].should.equal(expected_events)
 
 
+@mock.patch.dict("os.environ", {"MOTO_ENABLE_ISO_REGIONS": "true"})
 @pytest.mark.parametrize(
     "test_region", ["us-west-2", "cn-northwest-1", "us-isob-east-1"]
 )

tests/test_sts/test_sts.py

@@ -5,6 +5,7 @@ import boto3
 from botocore.client import ClientError
 from datetime import datetime
 from freezegun import freeze_time
+from unittest.mock import patch
 import pytest
 import sure  # noqa # pylint: disable=unused-import
 
@@ -707,6 +708,7 @@ def test_federation_token_with_too_long_policy():
     )
 
 
+@patch.dict("os.environ", {"MOTO_ENABLE_ISO_REGIONS": "true"})
 @pytest.mark.parametrize("region", ["us-west-2", "cn-northwest-1", "us-isob-east-1"])
 @mock_sts
 def test_sts_regions(region):