Improve recording decorators (#44180)

Author: scbedd

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

@@ -20,7 +20,7 @@
 from .exceptions import AzureTestError, ReservedResourceNameError
 from .proxy_fixtures import environment_variables, recorded_test, variable_recorder
 from .proxy_startup import start_test_proxy, stop_test_proxy, test_proxy
-from .proxy_testcase import recorded_by_proxy
+from .proxy_testcase import recorded_by_proxy, RecordedTransport
 from .sanitizers import (
     add_api_version_transform,
     add_batch_sanitizers,
@@ -105,6 +105,7 @@
     "EnvironmentVariableLoader",
     "environment_variables",
     "recorded_by_proxy",
+    "RecordedTransport",
     "recorded_test",
     "test_proxy",
     "trim_kwargs_from_test_function",

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

@@ -5,99 +5,183 @@
 # --------------------------------------------------------------------------
 import logging
 import urllib.parse as url_parse
+from functools import wraps
 
 from azure.core.exceptions import ResourceNotFoundError
 from azure.core.pipeline.policies import ContentDecodePolicy
 from azure.core.pipeline.transport import AioHttpTransport
 
+try:
+    import httpx
+
+    AsyncHTTPXTransport = httpx.AsyncHTTPTransport
+except ImportError:
+    httpx = None
+    AsyncHTTPXTransport = None
+
 from ..helpers import is_live_and_not_recording, trim_kwargs_from_test_function
 from ..proxy_testcase import (
+    RecordedTransport,
+    _transform_args,
+    _transform_httpx_args,
     get_test_id,
     start_record_or_playback,
-    transform_request,
+    restore_httpx_response_url,
     stop_record_or_playback,
 )
 
 
-def recorded_by_proxy_async(test_func):
-    """Decorator that redirects network requests to target the azure-sdk-tools test proxy. Use with recorded tests.
-
-    For more details and usage examples, refer to
-    https://github.com/Azure/azure-sdk-for-python/blob/main/doc/dev/tests.md#write-or-run-tests
+def recorded_by_proxy_async(*transports):
     """
+    Decorator for recording and playing back test proxy sessions in async tests.
 
-    async def record_wrap(*args, **kwargs):
-        def transform_args(*args, **kwargs):
-            copied_positional_args = list(args)
-            request = copied_positional_args[1]
+    Args:
+        *transports: Which transport(s) to record. Pass one or more comma separated RecordedTransport enum values.
+            - No args (default): Record AioHttpTransport.send calls (azure.core).
+            - RecordedTransport.AZURE_CORE: Record AioHttpTransport.send calls. Same as the default above.
+            - RecordedTransport.HTTPX: Record AsyncHTTPXTransport.handle_async_request calls.
+            - RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX: Record both transports.
 
-            transform_request(request, recording_id)
+    Usages:
 
-            return tuple(copied_positional_args), kwargs
+      from devtools_testutils.aio import recorded_by_proxy_async
+      from devtools_testutils import RecordedTransport
 
-        trimmed_kwargs = {k: v for k, v in kwargs.items()}
-        trim_kwargs_from_test_function(test_func, trimmed_kwargs)
+      # If your test uses azure.core only network calls (default)
+      @recorded_by_proxy_async
+      async def test(...): ...
 
-        if is_live_and_not_recording():
-            return await test_func(*args, **trimmed_kwargs)
+      # Explicitly enable azure.core recordings only (equivalent to the above)
+      @recorded_by_proxy_async(RecordedTransport.AZURE_CORE)
+      async def test(...): ...
 
-        test_id = get_test_id()
-        recording_id, variables = start_record_or_playback(test_id)
-        original_transport_func = AioHttpTransport.send
+      # If your test uses httpx only for network calls
+      @recorded_by_proxy_async(RecordedTransport.HTTPX)
+      async def test(...): ...
 
-        async def combined_call(*args, **kwargs):
-            adjusted_args, adjusted_kwargs = transform_args(*args, **kwargs)
-            result = await original_transport_func(*adjusted_args, **adjusted_kwargs)
-
-            # make the x-recording-upstream-base-uri the URL of the request
-            # this makes the request look like it was made to the original endpoint instead of to the proxy
-            # without this, things like LROPollers can get broken by polling the wrong endpoint
-            parsed_result = url_parse.urlparse(result.request.url)
-            upstream_uri = url_parse.urlparse(result.request.headers["x-recording-upstream-base-uri"])
-            upstream_uri_dict = {
-                "scheme": upstream_uri.scheme,
-                "netloc": upstream_uri.netloc,
-            }
-            original_target = parsed_result._replace(**upstream_uri_dict).geturl()
-
-            result.request.url = original_target
-            return result
+      # If your test uses both azure.core and httpx for network calls
+      @recorded_by_proxy_async(RecordedTransport.AZURE_CORE, RecordedTransport.HTTPX)
+      async def test(...): ...
+    """
 
-        AioHttpTransport.send = combined_call
+    # Bare decorator usage: @recorded_by_proxy_async
+    if len(transports) == 1 and callable(transports[0]):
+        test_func = transports[0]
+        transport_list = [(AioHttpTransport, "send")]
+        return _make_proxy_decorator_async(transport_list)(test_func)
+
+    # Parameterized decorator usage: @recorded_by_proxy_async(...)
+    # Determine which transports to use
+    transport_list = []
+
+    # If no transports specified, default to azure.core
+    transport_set = set(transports) if transports else {RecordedTransport.AZURE_CORE}
+
+    # Add transports based on what's in the set
+    for transport in transport_set:
+        if transport == RecordedTransport.AZURE_CORE or (
+            isinstance(transport, str) and transport == RecordedTransport.AZURE_CORE.value
+        ):
+            transport_list.append((AioHttpTransport, "send"))
+        elif transport == RecordedTransport.HTTPX or (
+            isinstance(transport, str) and transport == RecordedTransport.HTTPX.value
+        ):
+            if AsyncHTTPXTransport is not None:
+                transport_list.append((AsyncHTTPXTransport, "handle_async_request"))
+
+    # If still no transports, fall back to azure.core
+    if not transport_list:
+        transport_list = [(AioHttpTransport, "send")]
+
+    # Return a decorator function that will be applied to the test function
+    return lambda test_func: _make_proxy_decorator_async(transport_list)(test_func)
+
+
+def _make_proxy_decorator_async(transports):
+    def _decorator(test_func):
+        @wraps(test_func)
+        async def record_wrap(*args, **kwargs):
+            # ---- your existing trimming/early-exit logic ----
+            trimmed_kwargs = {k: v for k, v in kwargs.items()}
+            trim_kwargs_from_test_function(test_func, trimmed_kwargs)
+
+            if is_live_and_not_recording():
+                return await test_func(*args, **trimmed_kwargs)
+
+            test_id = get_test_id()
+            recording_id, variables = start_record_or_playback(test_id)
+
+            # Build a wrapper factory so each patched method closes over its own original
+            def make_combined_call(original_transport_func, is_httpx=False):
+                async def combined_call(*call_args, **call_kwargs):
+                    if is_httpx:
+                        adjusted_args, adjusted_kwargs = _transform_httpx_args(recording_id, *call_args, **call_kwargs)
+                        result = await original_transport_func(*adjusted_args, **adjusted_kwargs)
+                        restore_httpx_response_url(result)
+                    else:
+                        adjusted_args, adjusted_kwargs = _transform_args(recording_id, *call_args, **call_kwargs)
+                        result = await original_transport_func(*adjusted_args, **adjusted_kwargs)
+                        # rewrite request.url to the original upstream for LROs, etc.
+                        parsed_result = url_parse.urlparse(result.request.url)
+                        upstream_uri = url_parse.urlparse(result.request.headers["x-recording-upstream-base-uri"])
+                        upstream_uri_dict = {"scheme": upstream_uri.scheme, "netloc": upstream_uri.netloc}
+                        original_target = parsed_result._replace(**upstream_uri_dict).geturl()
+                        result.request.url = original_target
+                    return result
+
+                return combined_call
+
+            # Patch multiple transports and ensure restoration
+            test_variables = None
+            test_run = False
+            originals = []
+            # monkeypatch all requested transports
+            for owner, name in transports:
+                original = getattr(owner, name)
+                # Check if this is an httpx transport by comparing with httpx transport classes
+                is_httpx_transport = (AsyncHTTPXTransport is not None and owner is AsyncHTTPXTransport) or (
+                    httpx is not None and owner.__module__.startswith("httpx")
+                )
+                setattr(owner, name, make_combined_call(original, is_httpx=is_httpx_transport))
+                originals.append((owner, name, original))
 
-        # call the modified function
-        # we define test_variables before invoking the test so the variable is defined in case of an exception
-        test_variables = None
-        # this tracks whether the test has been run yet; used when calling the test function with/without `variables`
-        # running without `variables` in the `except` block leads to unnecessary exceptions in test execution output
-        test_run = False
-        try:
             try:
-                test_variables = await test_func(*args, variables=variables, **trimmed_kwargs)
-                test_run = True
-            except TypeError as error:
-                if "unexpected keyword argument" in str(error) and "variables" in str(error):
-                    logger = logging.getLogger()
-                    logger.info(
-                        "This test can't accept variables as input. The test method should accept `**kwargs` and/or a "
-                        "`variables` parameter to make use of recorded test variables."
-                    )
-                else:
-                    raise error
-            # if the test couldn't accept `variables`, run the test without passing them
-            if not test_run:
-                test_variables = await test_func(*args, **trimmed_kwargs)
-
-        except ResourceNotFoundError as error:
-            error_body = ContentDecodePolicy.deserialize_from_http_generics(error.response)
-            message = error_body.get("message") or error_body.get("Message")
-            error_with_message = ResourceNotFoundError(message=message, response=error.response)
-            raise error_with_message from error
-
-        finally:
-            AioHttpTransport.send = original_transport_func
-            stop_record_or_playback(test_id, recording_id, test_variables)
-
-        return test_variables
-
-    return record_wrap
+                try:
+                    test_variables = await test_func(*args, variables=variables, **trimmed_kwargs)
+                    test_run = True
+                except TypeError as error:
+                    if "unexpected keyword argument" in str(error) and "variables" in str(error):
+                        logger = logging.getLogger()
+                        logger.info(
+                            "This test can't accept variables as input. "
+                            "Accept `**kwargs` and/or a `variables` parameter to use recorded variables."
+                        )
+                    else:
+                        raise
+
+                if not test_run:
+                    test_variables = await test_func(*args, **trimmed_kwargs)
+
+            except ResourceNotFoundError as error:
+                error_body = ContentDecodePolicy.deserialize_from_http_generics(error.response)
+                troubleshoot = (
+                    "Playback failure -- for help resolving, see https://aka.ms/azsdk/python/test-proxy/troubleshoot."
+                )
+                message = error_body.get("message") or error_body.get("Message")
+                error_with_message = ResourceNotFoundError(
+                    message=f"{troubleshoot} Error details:\n{message}",
+                    response=error.response,
+                )
+                raise error_with_message from error
+
+            finally:
+                # restore in reverse order
+                for owner, name, original in reversed(originals):
+                    setattr(owner, name, original)
+                stop_record_or_playback(test_id, recording_id, test_variables)
+
+            return test_variables
+
+        return record_wrap
+
+    return _decorator