[Test Proxy] Make add_sanitizer a module-level method (Azure#20701)

zihzhan-msft · Sep 16, 2021 · 0645f49 · 0645f49
1 parent ddf49b1
commit 0645f49
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 14 deletions.
diff --git a/tools/azure-sdk-tools/devtools_testutils/__init__.py b/tools/azure-sdk-tools/devtools_testutils/__init__.py
@@ -1,5 +1,5 @@
 from .mgmt_testcase import AzureMgmtTestCase, AzureMgmtPreparer
-from .azure_recorded_testcase import AzureRecordedTestCase
+from .azure_recorded_testcase import add_sanitizer, AzureRecordedTestCase
 from .azure_testcase import AzureTestCase, is_live, get_region_override
 from .resource_testcase import (
     FakeResource,
@@ -21,6 +21,7 @@
 from .fake_credential import FakeTokenCredential
 
 __all__ = [
+    "add_sanitizer",
     "AzureMgmtTestCase",
     "AzureMgmtPreparer",
     "AzureRecordedTestCase",

diff --git a/tools/azure-sdk-tools/devtools_testutils/azure_recorded_testcase.py b/tools/azure-sdk-tools/devtools_testutils/azure_recorded_testcase.py
@@ -31,12 +31,42 @@
     pass
 
 if TYPE_CHECKING:
-    from typing import Optional
+    from typing import Any
 
 
 load_dotenv(find_dotenv())
 
 
+def add_sanitizer(sanitizer, **kwargs):
+    # type: (ProxyRecordingSanitizer, **Any) -> None
+    """Registers a sanitizer, matcher, or transform with the test proxy.
+
+    :param sanitizer: The name of the sanitizer, matcher, or transform you want to add.
+    :type sanitizer: ProxyRecordingSanitizer or str
+
+    :keyword str value: The substitution value.
+    :keyword str regex: A regex for a sanitizer. Can be defined as a simple regex, or if a ``group_for_replace`` is
+        provided, a substitution operation.
+    :keyword str group_for_replace: The capture group that needs to be operated upon. Do not provide if you're invoking
+        a simple replacement operation.
+    """
+    request_args = {}
+    request_args["value"] = kwargs.get("value") or "fakevalue"
+    request_args["regex"] = kwargs.get("regex") or "[a-z]+(?=(?:-secondary)\\.(?:table|blob|queue)\\.core\\.windows\\.net)"
+    request_args["group_for_replace"] = kwargs.get("group_for_replace")
+
+    if sanitizer == ProxyRecordingSanitizer.URI:
+        requests.post(
+            "{}/Admin/AddSanitizer".format(PROXY_URL),
+            headers={"x-abstraction-identifier": ProxyRecordingSanitizer.URI.value},
+            json={
+                "regex": request_args["regex"],
+                "value": request_args["value"],
+                "groupForReplace": request_args["group_for_replace"]
+            },
+        )
+
+
 def is_live():
     """A module version of is_live, that could be used in pytest marker."""
     if not hasattr(is_live, "_cache"):
@@ -81,18 +111,6 @@ def in_recording(self):
     def recording_processors(self):
         return []
 
-    def add_sanitizer(self, sanitizer, regex=None, value=None):
-        # type: (ProxyRecordingSanitizer, Optional[str], Optional[str]) -> None
-        if sanitizer == ProxyRecordingSanitizer.URI:
-            requests.post(
-                "{}/Admin/AddSanitizer".format(PROXY_URL),
-                headers={"x-abstraction-identifier": ProxyRecordingSanitizer.URI.value},
-                json={
-                    "regex": regex or "[a-z]+(?=(?:-secondary)\\.(?:table|blob|queue)\\.core\\.windows\\.net)",
-                    "value": value or "fakevalue"
-                },
-            )
-
     def is_playback(self):
         return not self.is_live