[Test Proxy] Add bodiless matcher, sanitize test SAS tokens (Azure#21823)

Author: mccoyp

tools/azure-sdk-tools/devtools_testutils/__init__.py

@@ -28,6 +28,7 @@
     add_remove_header_sanitizer,
     add_request_subscription_id_sanitizer,
     add_uri_regex_sanitizer,
+    set_bodiless_matcher,
 )
 from .helpers import ResponseCallback, RetryCounter
 from .fake_credential import FakeTokenCredential, ACCOUNT_FAKE_KEY
@@ -61,6 +62,7 @@
     "PowerShellPreparer",
     "recorded_by_proxy",
     "test_proxy",
+    "set_bodiless_matcher",
     "start_test_proxy",
     "stop_test_proxy",
     "ResponseCallback",

tools/azure-sdk-tools/devtools_testutils/azure_recorded_testcase.py

@@ -4,7 +4,6 @@
 # license information.
 # --------------------------------------------------------------------------
 import functools
-import logging
 import os
 import os.path
 import six
@@ -14,11 +13,12 @@
 
 from dotenv import load_dotenv, find_dotenv
 
-from azure_devtools.scenario_tests.config import TestConfig
 from azure_devtools.scenario_tests.utilities import trim_kwargs_from_test_function
 
 from . import mgmt_settings_fake as fake_settings
 from .azure_testcase import _is_autorest_v3, get_resource_name, get_qualified_method_name
+from .helpers import is_live
+from .sanitizers import add_general_regex_sanitizer
 
 try:
     # Try to import the AsyncFakeCredential, if we cannot assume it is Python 2
@@ -33,11 +33,12 @@
 load_dotenv(find_dotenv())
 
 
-def is_live():
-    """A module version of is_live, that could be used in pytest marker."""
-    if not hasattr(is_live, "_cache"):
-        is_live._cache = TestConfig().record_mode
-    return is_live._cache
+def _sanitize_token(token, fake_token):
+    add_general_regex_sanitizer(value=fake_token, regex=token)
+    url_safe_token = token.replace("/", u"%2F")
+    add_general_regex_sanitizer(value=fake_token, regex=url_safe_token)
+    async_token = token.replace(u"%3A", ":")
+    add_general_regex_sanitizer(value=fake_token, regex=async_token)
 
 
 class AzureRecordedTestCase(object):
@@ -204,6 +205,7 @@ def generate_sas(self, *args, **kwargs):
         token = sas_func(*sas_func_pos_args, **kwargs)
 
         fake_token = self._create_fake_token(token, fake_value)
+        _sanitize_token(token, fake_token)
 
         if self.is_live:
             return token

tools/azure-sdk-tools/devtools_testutils/helpers.py

@@ -3,6 +3,16 @@
 # Licensed under the MIT License. See License.txt in the project root for
 # license information.
 # --------------------------------------------------------------------------
+from azure_devtools.scenario_tests.config import TestConfig
+
+
+def is_live():
+    """A module version of is_live, that could be used in pytest marker."""
+    if not hasattr(is_live, "_cache"):
+        is_live._cache = TestConfig().record_mode
+    return is_live._cache
+
+
 class RetryCounter(object):
     def __init__(self):
         self.count = 0

tools/azure-sdk-tools/devtools_testutils/proxy_testcase.py

@@ -7,7 +7,9 @@
 import logging
 import requests
 import six
+import sys
 from typing import TYPE_CHECKING
+import pdb
 
 try:
     # py3
@@ -23,13 +25,12 @@
 
 # the trimming function to clean up incoming arguments to the test function we are wrapping
 from azure_devtools.scenario_tests.utilities import trim_kwargs_from_test_function
-from .azure_recorded_testcase import is_live
+from .helpers import is_live
 from .config import PROXY_URL
 
 if TYPE_CHECKING:
     from typing import Tuple
 
-
 # To learn about how to migrate SDK tests to the test proxy, please refer to the migration guide at
 # https://github.com/Azure/azure-sdk-for-python/blob/main/doc/dev/test_proxy_migration_guide.md
 
@@ -40,9 +41,15 @@
 PLAYBACK_START_URL = "{}/playback/start".format(PROXY_URL)
 PLAYBACK_STOP_URL = "{}/playback/stop".format(PROXY_URL)
 
-# TODO, create a pytest scope="session" implementation that can be added to a fixture such that unit tests can
-# startup/shutdown the local test proxy
-# this should also fire the admin mapping updates, and start/end the session for commiting recording updates
+# we store recording IDs in a module-level variable so that sanitizers can access them
+# we map test IDs to recording IDs, rather than storing only the current test's recording ID, for parallelization
+this = sys.modules[__name__]
+this.recording_ids = {}
+
+
+def get_recording_id():
+    test_id = get_test_id()
+    return this.recording_ids.get(test_id)
 
 
 def get_test_id():
@@ -92,6 +99,8 @@ def start_record_or_playback(test_id):
                     ValueError("The response body returned from starting playback did not contain valid JSON"), ex
                 )
 
+    # set recording ID in a module-level variable so that sanitizers can access it
+    this.recording_ids[test_id] = recording_id
     return (recording_id, variables)
 
 

tools/azure-sdk-tools/devtools_testutils/sanitizers.py

@@ -7,11 +7,23 @@
 from typing import TYPE_CHECKING
 
 from .config import PROXY_URL
+from .proxy_testcase import get_recording_id
 
 if TYPE_CHECKING:
     from typing import Any, Dict
 
 
+def set_bodiless_matcher():
+    # type: () -> None
+    """Adjusts the "match" operation to EXCLUDE the body when matching a request to a recording's entries.
+
+    This method must be called during test case execution, rather than at a session, module, or class level.
+    """
+
+    x_recording_id = get_recording_id()
+    _send_matcher_request("BodilessMatcher", {"x-recording-id": x_recording_id})
+
+
 def add_body_key_sanitizer(**kwargs):
     # type: (**Any) -> None
     """Registers a sanitizer that offers regex update of a specific JTokenPath within a returned body.
@@ -28,7 +40,7 @@ def add_body_key_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("BodyKeySanitizer", request_args)
+    _send_sanitizer_request("BodyKeySanitizer", request_args)
 
 
 def add_body_regex_sanitizer(**kwargs):
@@ -46,7 +58,7 @@ def add_body_regex_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("BodyRegexSanitizer", request_args)
+    _send_sanitizer_request("BodyRegexSanitizer", request_args)
 
 
 def add_continuation_sanitizer(**kwargs):
@@ -65,7 +77,7 @@ def add_continuation_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("ContinuationSanitizer", request_args)
+    _send_sanitizer_request("ContinuationSanitizer", request_args)
 
 
 def add_general_regex_sanitizer(**kwargs):
@@ -82,7 +94,7 @@ def add_general_regex_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("GeneralRegexSanitizer", request_args)
+    _send_sanitizer_request("GeneralRegexSanitizer", request_args)
 
 
 def add_header_regex_sanitizer(**kwargs):
@@ -102,14 +114,14 @@ def add_header_regex_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("HeaderRegexSanitizer", request_args)
+    _send_sanitizer_request("HeaderRegexSanitizer", request_args)
 
 
 def add_oauth_response_sanitizer():
     # type: () -> None
     """Registers a sanitizer that cleans out all request/response pairs that match an oauth regex in their URI."""
 
-    _send_request("OAuthResponseSanitizer", {})
+    _send_sanitizer_request("OAuthResponseSanitizer", {})
 
 
 def add_remove_header_sanitizer(**kwargs):
@@ -121,7 +133,7 @@ def add_remove_header_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("RemoveHeaderSanitizer", request_args)
+    _send_sanitizer_request("RemoveHeaderSanitizer", request_args)
 
 
 def add_request_subscription_id_sanitizer(**kwargs):
@@ -134,7 +146,7 @@ def add_request_subscription_id_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("ReplaceRequestSubscriptionId", request_args)
+    _send_sanitizer_request("ReplaceRequestSubscriptionId", request_args)
 
 
 def add_uri_regex_sanitizer(**kwargs):
@@ -149,35 +161,59 @@ def add_uri_regex_sanitizer(**kwargs):
     """
 
     request_args = _get_request_args(**kwargs)
-    _send_request("UriRegexSanitizer", request_args)
+    _send_sanitizer_request("UriRegexSanitizer", request_args)
 
 
 def _get_request_args(**kwargs):
     # type: (**Any) -> Dict
     """Returns a dictionary of sanitizer constructor headers"""
 
     request_args = {}
-    request_args["groupForReplace"] = kwargs.get("group_for_replace")
-    request_args["headersForRemoval"] = kwargs.get("headers")
-    request_args["jsonPath"] = kwargs.get("json_path")
-    request_args["key"] = kwargs.get("key")
-    request_args["method"] = kwargs.get("method")
-    request_args["regex"] = kwargs.get("regex")
-    request_args["resetAfterFirst"] = kwargs.get("reset_after_first")
-    request_args["value"] = kwargs.get("value")
+    if "group_for_replace" in kwargs:
+        request_args["groupForReplace"] = kwargs.get("group_for_replace")
+    if "headers" in kwargs:
+        request_args["headersForRemoval"] = kwargs.get("headers")
+    if "json_path" in kwargs:
+        request_args["jsonPath"] = kwargs.get("json_path")
+    if "key" in kwargs:
+        request_args["key"] = kwargs.get("key")
+    if "method" in kwargs:
+        request_args["method"] = kwargs.get("method")
+    if "regex" in kwargs:
+        request_args["regex"] = kwargs.get("regex")
+    if "reset_after_first" in kwargs:
+        request_args["resetAfterFirst"] = kwargs.get("reset_after_first")
+    if "value" in kwargs:
+        request_args["value"] = kwargs.get("value")
     return request_args
 
 
-def _send_request(sanitizer, parameters):
+def _send_matcher_request(matcher, headers):
+    # type: (str, Dict) -> None
+    """Sends a POST request to the test proxy endpoint to register the specified matcher.
+
+    :param str matcher: The name of the matcher to set.
+    :param dict headers: Any matcher headers, as a dictionary.
+    """
+
+    headers_to_send = {"x-abstraction-identifier": matcher}
+    headers_to_send.update(headers)
+    requests.post(
+        "{}/Admin/SetMatcher".format(PROXY_URL),
+        headers=headers_to_send,
+    )
+
+
+def _send_sanitizer_request(sanitizer, parameters):
     # type: (str, Dict) -> None
-    """Send a POST request to the test proxy endpoint to register the specified sanitizer.
+    """Sends a POST request to the test proxy endpoint to register the specified sanitizer.
 
     :param str sanitizer: The name of the sanitizer to add.
     :param dict parameters: The sanitizer constructor parameters, as a dictionary.
     """
 
     requests.post(
         "{}/Admin/AddSanitizer".format(PROXY_URL),
-        headers={"x-abstraction-identifier": sanitizer},
+        headers={"x-abstraction-identifier": sanitizer, "Content-Type": "application/json"},
         json=parameters
     )