[Test Proxy] Support AsyncioRequestsTransport
@@ -8,7 +8,7 @@ import urllib.parse as url_parse
from azure.core.exceptions import ResourceNotFoundError
from azure.core.pipeline.policies import ContentDecodePolicy
-from azure.core.pipeline.transport import AioHttpTransport
+from azure.core.pipeline.transport import AioHttpTransport, AsyncioRequestsTransport
from azure_devtools.scenario_tests.utilities import trim_kwargs_from_test_function
from ..helpers import is_live_and_not_recording
@@ -44,11 +44,15 @@ def recorded_by_proxy_async(test_func):
test_id = get_test_id()
recording_id, variables = start_record_or_playback(test_id)
- original_transport_func = AioHttpTransport.send
+ original_aio_transport_func = AioHttpTransport.send
+ original_asyncio_transport_func = AsyncioRequestsTransport.send
async def combined_call(*args, **kwargs):
adjusted_args, adjusted_kwargs = transform_args(*args, **kwargs)
- result = await original_transport_func(*adjusted_args, **adjusted_kwargs)
+ try:
+ result = await original_aio_transport_func(*adjusted_args, **adjusted_kwargs)
+ except TypeError: # if awaiting AioHttpTransport.send fails, it means we're using AsyncioRequestsTransport
+ result = await original_asyncio_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
@@ -65,6 +69,7 @@ def recorded_by_proxy_async(test_func):
return result
AioHttpTransport.send = combined_call
+ AsyncioRequestsTransport.send = combined_call
# call the modified function
# we define test_output before invoking the test so the variable is defined in case of an exception
@@ -85,7 +90,8 @@ def recorded_by_proxy_async(test_func):
error_with_message = ResourceNotFoundError(message=message, response=error.response)
raise error_with_message from error
finally:
- AioHttpTransport.send = original_transport_func
+ AioHttpTransport.send = original_aio_transport_func
+ AsyncioRequestsTransport.send = original_asyncio_transport_func
stop_record_or_playback(test_id, recording_id, test_output)
return test_output
