From 0645f499eb0aaf253f45fb5146f8f456fd9baa25 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?McCoy=20Pati=C3=B1o?=
 <39780829+mccoyp@users.noreply.github.com>
Date: Thu, 16 Sep 2021 15:58:10 -0700
Subject: [PATCH] [Test Proxy] Make add_sanitizer a module-level method
 (#20701)

---
 .../devtools_testutils/__init__.py            |  3 +-
 .../azure_recorded_testcase.py                | 44 +++++++++++++------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/tools/azure-sdk-tools/devtools_testutils/__init__.py b/tools/azure-sdk-tools/devtools_testutils/__init__.py
index 69db6f41dfb3..c58ea74788c2 100644
--- 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
index d39cfe79960c..96452251da3d 100644
--- 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