[async] test fixes:

Fixed synch factory in async code of _wif.py
Fixed async make_session and added necessary host_port_pooling

Author: sfc-gh-fpawlowski

src/snowflake/connector/aio/_session_manager.py

@@ -540,17 +540,24 @@ def make_session(self, *, url: str | None = None) -> aiohttp.ClientSession:
             session_manager=self.clone(),
             snowflake_ocsp_mode=self._cfg.snowflake_ocsp_mode,
         )
-        # We use requests.utils here (in asynch code) to keep the behaviour uniform for synch and asynch code. If we wanted each version to depict its http library's behaviour, we could use here: aiohttp.helpers.proxy_bypass(url, proxies={...}) here
-        proxy = (
-            None
-            if should_bypass_proxies(url, no_proxy=self.config.no_proxy)
-            else self.proxy_url
-        )
+
+        proxy_from_conn_params: str | None = None
+        if not aiohttp.helpers.proxies_from_env():
+            # TODO: This is only needed because we want to keep compatibility with the synch driver version.
+            #   Otherwise, we could remove that condition and always pass proxy from conn params to the Session constructor.
+            #   But in such case precedence will be reverted and it will overwrite the env vars settings.
+
+            # We use requests.utils here (in asynch code) to keep the behaviour uniform for synch and asynch code. If we wanted each version to depict its http library's behaviour, we could use here: aiohttp.helpers.proxy_bypass(url, proxies={...}) here
+            proxy_from_conn_params = (
+                None
+                if should_bypass_proxies(url, no_proxy=self.config.no_proxy)
+                else self.proxy_url
+            )
         # Construct session with base proxy set, request() may override per-URL when bypassing
         return self.SessionWithProxy(
             connector=connector,
             trust_env=self._cfg.trust_env,
-            proxy=proxy,
+            proxy=proxy_from_conn_params,
         )
 
 

src/snowflake/connector/aio/_wif_util.py

@@ -14,7 +14,6 @@
 
 from ..errorcode import ER_WIF_CREDENTIALS_NOT_FOUND
 from ..errors import MissingDependencyError, ProgrammingError
-from ..session_manager import SessionManagerFactory
 from ..wif_util import (
     DEFAULT_ENTRA_SNOWFLAKE_RESOURCE,
     SNOWFLAKE_AUDIENCE,
@@ -24,7 +23,7 @@
     extract_iss_and_sub_without_signature_verification,
     get_aws_sts_hostname,
 )
-from ._session_manager import SessionManager
+from ._session_manager import SessionManager, SessionManagerFactory
 
 logger = logging.getLogger(__name__)
 

test/unit/aio/mock_utils.py

@@ -35,7 +35,7 @@ async def forbidden_connect(*args, **kwargs):
 
     class MockSessionManager(SessionManager):
         def make_session(self, *, url: str | None = None):
-            session = super().make_session(url)
+            session = super().make_session(url=url)
             if not allow_send:
                 # Block at connector._connect level (like sync blocks session.send)
                 # This allows patches on session.request to work

test/unit/aio/test_connection_async_unit.py

@@ -912,7 +912,7 @@ def test_connect_metadata_preservation():
         connect_doc == source_doc
     ), "inspect.getdoc(connect) should match inspect.getdoc(SnowflakeConnection.__init__)"
 
-    # Test 8: Check that connect is callable and returns expected type
+    # Test 8: Check that connect is callable
     assert callable(connect), "connect should be callable"
 
     # Test 9: Check type() and __class__ values (important for user introspection)
@@ -933,3 +933,4 @@ def test_connect_metadata_preservation():
     assert (
         len(params) > 0
     ), "connect should have parameters from SnowflakeConnection.__init__"
+    # Should have parameters like account, user, password, etc.

test/unit/aio/test_proxies_async.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import urllib.request
 from collections import deque
 from test.unit.test_proxies import (
     DbRequestFlags,
@@ -15,6 +16,9 @@
 
 import aiohttp
 import pytest
+from aiohttp import BasicAuth
+from aiohttp.helpers import proxies_from_env
+from yarl import URL
 
 from snowflake.connector.aio import connect as async_connect
 
@@ -355,6 +359,57 @@ async def test_no_proxy_basic_param_proxy_bypass_backend(
     assert flags.proxy_saw_storage is True
 
 
+@pytest.fixture
+def fix_aiohttp_proxy_bypass(monkeypatch):
+    """Fix aiohttp's proxy bypass to check host:port instead of just host.
+
+    This fixture implements a two-step fix:
+    1. Override get_env_proxy_for_url to use host_port_subcomponent for proxy_bypass
+    2. Override urllib.request._splitport to return (host:port, port) for proper matching
+    """
+
+    # Step 1: Override get_env_proxy_for_url to pass host:port to proxy_bypass
+    def get_env_proxy_for_url_with_port(url: URL) -> tuple[URL, BasicAuth | None]:
+        """Get a permitted proxy for the given URL from the env, checking host:port."""
+        from urllib.request import proxy_bypass
+
+        # Check proxy bypass using host:port combination
+        if url.host is not None:
+            # Use host_port_subcomponent which includes port
+            host_port = f"{url.host}:{url.port}" if url.port else url.host
+            if proxy_bypass(host_port):
+                raise LookupError(f"Proxying is disallowed for `{host_port!r}`")
+
+        proxies_in_env = proxies_from_env()
+        try:
+            proxy_info = proxies_in_env[url.scheme]
+        except KeyError:
+            raise LookupError(f"No proxies found for `{url!s}` in the env")
+        else:
+            return proxy_info.proxy, proxy_info.proxy_auth
+
+    # Step 2: Override _splitport to return host:port as first element
+    original_splitport = urllib.request._splitport
+
+    def _splitport_with_port(host):
+        """Override to return (host:port, port) instead of (host, port)."""
+        result = original_splitport(host)
+        if result is None:
+            return (host, None)
+        host_only, port = result
+        # If port was found, return the original host (with port) as first element
+        if port is not None:
+            return (host, port)  # Return original host:port string
+        return (host_only, port)
+
+    monkeypatch.setattr(
+        aiohttp.client, "get_env_proxy_for_url", get_env_proxy_for_url_with_port
+    )
+    monkeypatch.setattr(urllib.request, "_splitport", _splitport_with_port)
+
+    yield
+
+
 @pytest.mark.skipolddriver
 @pytest.mark.parametrize("proxy_method", ["explicit_args", "env_vars"])
 @pytest.mark.parametrize("no_proxy_source", ["param", "env"])
@@ -366,6 +421,7 @@ async def test_no_proxy_source_vs_proxy_method_matrix(
     proxy_method,
     no_proxy_source,
     host_port_pooling,
+    fix_aiohttp_proxy_bypass,
 ):
     if proxy_method == "env_vars" and no_proxy_source == "param":
         pytest.xfail(
@@ -407,6 +463,7 @@ async def test_no_proxy_backend_matrix(
     proxy_method,
     no_proxy_source,
     host_port_pooling,
+    fix_aiohttp_proxy_bypass,
 ):
     if proxy_method == "env_vars" and no_proxy_source == "param":
         pytest.xfail(
@@ -461,6 +518,7 @@ async def test_no_proxy_multiple_values_param_only(
     wiremock_mapping_dir,
     proxy_env_vars,
     no_proxy_factory,
+    host_port_pooling,  # Unlike in synch code - Session stores no_proxy setup so it would be reused for proxy and backend since they are both on localhost
 ):
     target_wm, storage_wm, proxy_wm = wiremock_backend_storage_proxy
     _setup_backend_storage_mappings(
@@ -569,6 +627,8 @@ async def test_proxy_env_vars_take_precedence_over_connection_params(
     wiremock_generic_mappings_dir,
     proxy_env_vars,
     monkeypatch,
+    host_port_pooling,
+    fix_aiohttp_proxy_bypass,
 ):
     """Verify that proxy_host/proxy_port connection parameters take precedence over env vars.