expose http_session_cls in AioConfig (#1102)

Author: jakkdl

CHANGES.rst

@@ -1,6 +1,10 @@
 Changes
 -------
 
+2.12.2 (2024-04-01)
+^^^^^^^^^^^^^^^^^^^
+* expose configuration of ``http_session_cls`` in ``AioConfig``
+
 2.12.1 (2024-03-04)
 ^^^^^^^^^^^^^^^^^^^
 * fix use of proxies #1070

aiobotocore/__init__.py

@@ -1 +1 @@
-__version__ = '2.12.1'
+__version__ = '2.12.2'

aiobotocore/args.py

@@ -5,7 +5,7 @@
 from botocore.args import ClientArgsCreator
 
 from .config import AioConfig
-from .endpoint import AioEndpointCreator
+from .endpoint import DEFAULT_HTTP_SESSION_CLS, AioEndpointCreator
 from .regions import AioEndpointRulesetResolver
 from .signers import AioRequestSigner
 
@@ -67,8 +67,10 @@ def get_client_args(
         # aiobotocore addition
         if isinstance(client_config, AioConfig):
             connector_args = client_config.connector_args
+            http_session_cls = client_config.http_session_cls
         else:
             connector_args = None
+            http_session_cls = DEFAULT_HTTP_SESSION_CLS
 
         new_config = AioConfig(connector_args, **config_kwargs)
         endpoint_creator = AioEndpointCreator(event_emitter)
@@ -79,9 +81,10 @@ def get_client_args(
             endpoint_url=endpoint_config['endpoint_url'],
             verify=verify,
             response_parser_factory=self._response_parser_factory,
+            timeout=(new_config.connect_timeout, new_config.read_timeout),
             max_pool_connections=new_config.max_pool_connections,
+            http_session_cls=http_session_cls,
             proxies=new_config.proxies,
-            timeout=(new_config.connect_timeout, new_config.read_timeout),
             socket_options=socket_options,
             client_cert=new_config.client_cert,
             proxies_config=new_config.proxies_config,

aiobotocore/config.py

@@ -3,13 +3,21 @@
 import botocore.client
 from botocore.exceptions import ParamValidationError
 
+from aiobotocore.endpoint import DEFAULT_HTTP_SESSION_CLS
+
 
 class AioConfig(botocore.client.Config):
-    def __init__(self, connector_args=None, **kwargs):
+    def __init__(
+        self,
+        connector_args=None,
+        http_session_cls=DEFAULT_HTTP_SESSION_CLS,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
 
         self._validate_connector_args(connector_args)
         self.connector_args = copy.copy(connector_args)
+        self.http_session_cls = http_session_cls
         if not self.connector_args:
             self.connector_args = dict()
 

aiobotocore/endpoint.py

@@ -19,6 +19,8 @@
 from aiobotocore.httpsession import AIOHTTPSession
 from aiobotocore.response import StreamingBody
 
+DEFAULT_HTTP_SESSION_CLS = AIOHTTPSession
+
 
 async def convert_to_response_dict(http_response, operation_model):
     """Convert an HTTP response object to a request dict.
@@ -295,7 +297,7 @@ def create_endpoint(
         response_parser_factory=None,
         timeout=DEFAULT_TIMEOUT,
         max_pool_connections=MAX_POOL_CONNECTIONS,
-        http_session_cls=AIOHTTPSession,
+        http_session_cls=DEFAULT_HTTP_SESSION_CLS,
         proxies=None,
         socket_options=None,
         client_cert=None,

tests/test_config.py

@@ -6,6 +6,7 @@
 from botocore.exceptions import ParamValidationError, ReadTimeoutError
 
 from aiobotocore.config import AioConfig
+from aiobotocore.httpsession import AIOHTTPSession
 from aiobotocore.session import AioSession, get_session
 from tests.mock_server import AIOServer
 
@@ -132,3 +133,27 @@ def test_merge():
     assert isinstance(new_config, AioConfig)
     assert new_config is not config
     assert new_config is not other_config
+
+
+# Check that it's possible to specify custom http_session_cls
+@pytest.mark.moto
+@pytest.mark.asyncio
+async def test_config_http_session_cls():
+    class SuccessExc(Exception):
+        ...
+
+    class MyHttpSession(AIOHTTPSession):
+        async def send(self, request):
+            raise SuccessExc
+
+    config = AioConfig(http_session_cls=MyHttpSession)
+    session = AioSession()
+    async with AIOServer() as server, session.create_client(
+        's3',
+        config=config,
+        endpoint_url=server.endpoint_url,
+        aws_secret_access_key='xxx',
+        aws_access_key_id='xxx',
+    ) as s3_client:
+        with pytest.raises(SuccessExc):
+            await s3_client.get_object(Bucket='foo', Key='bar')