Fixes #25

docker_registry_client_async/dockerregistryclientasync.py

@@ -84,26 +84,34 @@ class DockerRegistryClientAsync:
     def __init__(
         self,
         *,
+        client_session: ClientSession = None,
+        client_session_kwargs: Dict = None,
         credentials_store: Path = None,
         no_proxy: str = None,
         proxies: Dict[str, str] = None,
         proxy_auth: BasicAuth = None,
         resolver_kwargs: Dict = None,
         ssl: Union[None, bool, Fingerprint, SSLContext] = None,
+        tcp_connector_kwargs: Dict = None,
         token_based_endpoints: List[str] = None,
         **kwargs,
     ):
         # pylint: disable=unused-argument
         """
         Args:
+            client_session: The underlying client session to use when making connections.
+            client_session_kwargs: Arguments to be passed to the client session.
             credentials_store: Path to the docker registry credentials store.
             no_proxy: A comma separated list of domains to exclude from proxying.
             proxies: Mapping of protocols to proxy urls, optionally including credentials.
             proxy_auth: The credentials to use when proxying.
             resolver_kwargs: Arguments to be passed to the resolver
             ssl: SSL context.
+            tcp_connector_kwargs: Arguments to be passed to the TCP connector.
             token_based_endpoints: List of token-based endpoints
         """
+        if not client_session_kwargs:
+            client_session_kwargs = {}
         if not credentials_store:
             credentials_store = Path(
                 os.environ.get(
@@ -135,20 +143,24 @@ def __init__(
         if ssl and DockerRegistryClientAsync.DEBUG:
             LOGGER.debug("SSL Context: %s", ssl.cert_store_stats())
 
+        if not tcp_connector_kwargs:
+            tcp_connector_kwargs = {}
         if not token_based_endpoints:
             token_based_endpoints = os.environ.get(
                 "DRCA_TOKEN_BASED_ENDPOINTS",
                 DockerRegistryClientAsync.DEFAULT_TOKEN_BASED_ENDPOINTS,
             ).split(",")
 
-        self.client_session = None
+        self.client_session = client_session
+        self.client_session_kwargs = client_session_kwargs
         self.credentials_store = credentials_store
         self.credentials = None
         self.proxies = proxies
         self.proxy_auth = proxy_auth
         self.proxy_no = no_proxy
         self.resolver_kwargs = resolver_kwargs
         self.ssl = ssl
+        self.tcp_connector_kwargs = tcp_connector_kwargs
         # Endpoint -> scope -> token
         self.tokens = {}
         self.token_based_endpoints = token_based_endpoints
@@ -246,11 +258,18 @@ async def _get_client_session(self) -> ClientSession:
             The AIOHTTP client session.
         """
         if not self.client_session:
-            self.client_session = ClientSession(
-                connector=TCPConnector(
-                    resolver=AsyncResolver(**self.resolver_kwargs), ssl=self.ssl
+            if "resolver" not in self.tcp_connector_kwargs:
+                self.tcp_connector_kwargs["resolver"] = AsyncResolver(
+                    **self.resolver_kwargs
                 )
-            )
+            if "ssl" not in self.tcp_connector_kwargs:
+                self.tcp_connector_kwargs["ssl"] = self.ssl
+            if "connector" not in self.client_session_kwargs:
+                self.client_session_kwargs["connector"] = TCPConnector(
+                    **self.tcp_connector_kwargs
+                )
+            self.client_session = ClientSession(**self.client_session_kwargs)
+
         return self.client_session
 
     async def _get_credentials(self, endpoint: str) -> Optional[str]:

tests/test_dockerregistryclientasync.py

@@ -8,6 +8,7 @@
 import hashlib
 import json
 import logging
+import socket
 
 from http import HTTPStatus
 from pathlib import Path
@@ -16,6 +17,7 @@
 import aiofiles
 import pytest
 
+from aiohttp import TCPConnector, ThreadedResolver
 from pytest_docker_registry_fixtures import (
     DockerRegistrySecure,
     ImageName as PDRFImageName,
@@ -2069,5 +2071,57 @@ async def test_put_manifest_from_disk_sync(
     assert response.digest == image_name.resolve_digest()
 
 
+@pytest.mark.online
+async def test_issue_25(
+    docker_registry_secure: DockerRegistrySecure, known_good_image: TypingKnownGoodImage
+):
+    """Test issue #25."""
+    credentials = docker_registry_secure.auth_header["Authorization"].split()[1]
+
+    image_names = list(get_identifier_map(known_good_image).keys())
+    image_name = image_names[0]
+
+    # IPv4 only ...
+    async with DockerRegistryClientAsync(
+        tcp_connector_kwargs={"family": socket.AF_INET},
+        ssl=docker_registry_secure.ssl_context,
+    ) as docker_registry_client_async:
+        await docker_registry_client_async.add_credentials(
+            credentials=credentials, endpoint=docker_registry_secure.endpoint
+        )
+        LOGGER.debug("Checking manifest for: %s ...", image_name)
+        response = await docker_registry_client_async.head_manifest(image_name)
+        assert response.client_response
+        assert response.result
+
+    # ThreadedResolver ...
+    async with DockerRegistryClientAsync(
+        tcp_connector_kwargs={"resolver": ThreadedResolver()},
+        ssl=docker_registry_secure.ssl_context,
+    ) as docker_registry_client_async:
+        await docker_registry_client_async.add_credentials(
+            credentials=credentials, endpoint=docker_registry_secure.endpoint
+        )
+        LOGGER.debug("Checking manifest for: %s ...", image_name)
+        response = await docker_registry_client_async.head_manifest(image_name)
+        assert response.client_response
+        assert response.result
+
+    # Custom TCPConnector ...
+    tcp_connector = TCPConnector(
+        resolver=ThreadedResolver(), ssl=docker_registry_secure.ssl_context
+    )
+    async with DockerRegistryClientAsync(
+        client_session_kwargs={"connector": tcp_connector}
+    ) as docker_registry_client_async:
+        await docker_registry_client_async.add_credentials(
+            credentials=credentials, endpoint=docker_registry_secure.endpoint
+        )
+        LOGGER.debug("Checking manifest for: %s ...", image_name)
+        response = await docker_registry_client_async.head_manifest(image_name)
+        assert response.client_response
+        assert response.result
+
+
 # TODO: Total image pull (with exists() checks)
 # TODO: Total image push