[Test Proxy] Add RecordedByProxy decorator and AzureRecordedTestCase (#…

…20138)

.gitignore

@@ -114,3 +114,6 @@ sdk/cosmos/azure-cosmos/test/test_config.py
 
 # env vars
 .env
+
+# local SSL certificate folder
+.certificate

sdk/tables/azure-data-tables/tests/test_table_service_stats.py

@@ -10,7 +10,7 @@
 from preparers import tables_decorator
 
 # --Test Class -----------------------------------------------------------------
-class TableServiceStatsTest(AzureTestCase, TableTestCase):
+class TestTableServiceStats(AzureTestCase, TableTestCase):
 
     # --Test cases per service ---------------------------------------
     @tables_decorator

sdk/tables/azure-data-tables/tests/test_table_service_stats_async.py

@@ -19,7 +19,7 @@
                          '></StorageServiceStats> '
 
 
-class TableServiceStatsTest(AzureTestCase, AsyncTableTestCase):
+class TestTableServiceStats(AzureTestCase, AsyncTableTestCase):
 
     @staticmethod
     def override_response_body_with_unavailable_status(response):

tools/azure-devtools/src/azure_devtools/scenario_tests/preparers.py

@@ -6,7 +6,6 @@
 import contextlib
 import functools
 import logging
-import sys
 from collections import namedtuple
 from threading import Lock
 
@@ -135,7 +134,14 @@ def _preparer_wrapper(test_class_instance, **kwargs):
                         )
 
             if test_class_instance.is_live:
-                test_class_instance.scrubber.register_name_pair(resource_name, self.moniker)
+                # Adding this for new proxy testcase
+                if hasattr(test_class_instance, "scrubber"):
+                    test_class_instance.scrubber.register_name_pair(resource_name, self.moniker)
+                else:
+                    _logger.info(
+                        "This test class instance has no scrubber, so the AbstractPreparer will not scrub any values "
+                        "in recordings."
+                    )
 
             # We shouldn't trim the same kwargs that we use for deletion,
             # we may remove some of the variables we needed to do the delete.

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

@@ -1,4 +1,5 @@
 from .mgmt_testcase import AzureMgmtTestCase, AzureMgmtPreparer
+from .azure_recorded_testcase import AzureRecordedTestCase
 from .azure_testcase import AzureTestCase, is_live, get_region_override
 from .resource_testcase import (
     FakeResource,
@@ -14,12 +15,15 @@
 )
 from .keyvault_preparer import KeyVaultPreparer
 from .powershell_preparer import PowerShellPreparer
+from .proxy_testcase import RecordedByProxy
+from .enums import ProxyRecordingSanitizer
 from .helpers import ResponseCallback, RetryCounter
 from .fake_credential import FakeTokenCredential
 
 __all__ = [
     "AzureMgmtTestCase",
     "AzureMgmtPreparer",
+    "AzureRecordedTestCase",
     "FakeResource",
     "ResourceGroupPreparer",
     "StorageAccountPreparer",
@@ -33,6 +37,8 @@
     "RandomNameResourceGroupPreparer",
     "CachedResourceGroupPreparer",
     "PowerShellPreparer",
+    "ProxyRecordingSanitizer",
+    "RecordedByProxy",
     "ResponseCallback",
     "RetryCounter",
     "FakeTokenCredential",

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

@@ -0,0 +1,3 @@
+from .proxy_testcase_async import RecordedByProxyAsync
+
+__all__ = ["RecordedByProxyAsync"]

tools/azure-sdk-tools/devtools_testutils/aio/proxy_testcase_async.py

@@ -0,0 +1,50 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+from azure.core.pipeline.transport import AioHttpTransport
+
+from azure_devtools.scenario_tests.utilities import trim_kwargs_from_test_function
+from ..proxy_testcase import (
+    get_test_id,
+    start_record_or_playback,
+    transform_request,
+    stop_record_or_playback,
+)
+
+
+def RecordedByProxyAsync(func):
+    async def record_wrap(*args, **kwargs):
+        test_id = get_test_id()
+        recording_id = start_record_or_playback(test_id)
+
+        def transform_args(*args, **kwargs):
+            copied_positional_args = list(args)
+            request = copied_positional_args[1]
+
+            transform_request(request, recording_id)
+
+            return tuple(copied_positional_args), kwargs
+
+        trimmed_kwargs = {k: v for k, v in kwargs.items()}
+        trim_kwargs_from_test_function(func, trimmed_kwargs)
+
+        original_func = AioHttpTransport.send
+
+        async def combined_call(*args, **kwargs):
+            adjusted_args, adjusted_kwargs = transform_args(*args, **kwargs)
+            return await original_func(*adjusted_args, **adjusted_kwargs)
+
+        AioHttpTransport.send = combined_call
+
+        # call the modified function.
+        try:
+            value = await func(*args, **trimmed_kwargs)
+        finally:
+            AioHttpTransport.send = original_func
+            stop_record_or_playback(test_id, recording_id)
+
+        return value
+
+    return record_wrap