add tests for _assign_requests_to_connections

Author: VictorPrins

tests/_async/test_connection_pool.py

@@ -892,3 +892,122 @@ async def test_keepalive_idle_connections():
             repr(pool)
             == "<AsyncConnectionPool [Requests: 0 active, 0 queued | Connections: 0 active, 1 idle]>"
         )
+
+
+@pytest.mark.anyio
+async def test_idle_but_not_available_connection_closing():
+    """
+    Test that a connection that is idle but not available gets closed.
+    This is a pathological edge case that shouldn't occur in reality, but the connection pool
+    handles it anyway.
+    """
+
+    class MockIdleNotAvailableConnection(httpcore.AsyncConnectionInterface):
+        def __init__(self) -> None:
+            self._origin = httpcore.Origin(b"https", b"example.com", 443)
+
+        def is_available(self) -> bool:
+            return False  # Not available
+
+        def is_idle(self) -> bool:
+            return True  # But is idle - this is the edge case
+
+        def is_closed(self) -> bool:
+            return False
+
+        def has_expired(self) -> bool:
+            return False
+
+    network_backend = httpcore.AsyncMockBackend([])
+
+    async with httpcore.AsyncConnectionPool(
+        network_backend=network_backend,
+    ) as pool:
+        # Replace the connection list with our pathological mock connection
+        mock_conn = MockIdleNotAvailableConnection()
+
+        pool._connections = [mock_conn]
+
+        # This should move the idle-but-not-available connection to closing_conns
+        closing_conns = pool._assign_requests_to_connections()
+
+        # Verify the connection is marked for closing
+        assert len(closing_conns) == 1
+        assert closing_conns[0] is mock_conn
+
+        # Verify it's no longer in the connection pool
+        assert len(pool._connections) == 0
+
+
+@pytest.mark.anyio
+async def test_active_http2_connection_keepalive_preservation():
+    """
+    Test that active HTTP/2 connections with capacity are preserved during keepalive enforcement.
+    This tests line 421 where active HTTP/2 connections are kept during keepalive enforcement.
+    """
+
+    class MockActiveHTTP2Connection(httpcore.AsyncConnectionInterface):
+        def __init__(self) -> None:
+            self._origin = httpcore.Origin(b"https", b"example.com", 443)
+
+        def is_available(self) -> bool:
+            return True  # Available with capacity
+
+        def is_idle(self) -> bool:
+            return False  # Not idle (active HTTP/2 connection with streams)
+
+        def is_closed(self) -> bool:
+            return False
+
+        def has_expired(self) -> bool:
+            return False
+
+        def get_available_stream_capacity(self) -> int:
+            return 5  # HTTP/2 connection with available stream capacity
+
+        async def aclose(self):
+            pass
+
+    class MockIdleConnection(httpcore.AsyncConnectionInterface):
+        def __init__(self) -> None:
+            self._origin = httpcore.Origin(b"https", b"example.com", 443)
+
+        def is_available(self) -> bool:
+            return True
+
+        def is_idle(self) -> bool:
+            return True  # This is an idle connection
+
+        def is_closed(self) -> bool:
+            return False
+
+        def has_expired(self) -> bool:
+            return False
+
+        def get_available_stream_capacity(self) -> int:
+            return 1
+
+    network_backend = httpcore.AsyncMockBackend([])
+
+    async with httpcore.AsyncConnectionPool(
+        max_keepalive_connections=0,  # Force keepalive limit to 0
+        http2=True,
+        network_backend=network_backend,
+    ) as pool:
+        # Create one idle connection and one active HTTP/2 connection
+        idle_conn = MockIdleConnection()
+        active_http2_conn = MockActiveHTTP2Connection()
+
+        with pool._optional_thread_lock:
+            pool._connections = [idle_conn, active_http2_conn]
+
+            # This should close idle connections but preserve the active HTTP/2 connection
+            closing_conns = pool._assign_requests_to_connections()
+
+            # Verify idle connection is marked for closing
+            assert len(closing_conns) == 1
+            assert idle_conn in closing_conns
+
+            # Verify the active HTTP/2 connection is preserved (line 421)
+            assert len(pool._connections) == 1
+            assert active_http2_conn in pool._connections

tests/_sync/test_connection_pool.py

@@ -892,3 +892,122 @@ def test_keepalive_idle_connections():
             repr(pool)
             == "<ConnectionPool [Requests: 0 active, 0 queued | Connections: 0 active, 1 idle]>"
         )
+
+
+
+def test_idle_but_not_available_connection_closing():
+    """
+    Test that a connection that is idle but not available gets closed.
+    This is a pathological edge case that shouldn't occur in reality, but the connection pool
+    handles it anyway.
+    """
+
+    class MockIdleNotAvailableConnection(httpcore.ConnectionInterface):
+        def __init__(self) -> None:
+            self._origin = httpcore.Origin(b"https", b"example.com", 443)
+
+        def is_available(self) -> bool:
+            return False  # Not available
+
+        def is_idle(self) -> bool:
+            return True  # But is idle - this is the edge case
+
+        def is_closed(self) -> bool:
+            return False
+
+        def has_expired(self) -> bool:
+            return False
+
+    network_backend = httpcore.MockBackend([])
+
+    with httpcore.ConnectionPool(
+        network_backend=network_backend,
+    ) as pool:
+        # Replace the connection list with our pathological mock connection
+        mock_conn = MockIdleNotAvailableConnection()
+
+        pool._connections = [mock_conn]
+
+        # This should move the idle-but-not-available connection to closing_conns
+        closing_conns = pool._assign_requests_to_connections()
+
+        # Verify the connection is marked for closing
+        assert len(closing_conns) == 1
+        assert closing_conns[0] is mock_conn
+
+        # Verify it's no longer in the connection pool
+        assert len(pool._connections) == 0
+
+
+
+def test_active_http2_connection_keepalive_preservation():
+    """
+    Test that active HTTP/2 connections with capacity are preserved during keepalive enforcement.
+    This tests line 421 where active HTTP/2 connections are kept during keepalive enforcement.
+    """
+
+    class MockActiveHTTP2Connection(httpcore.ConnectionInterface):
+        def __init__(self) -> None:
+            self._origin = httpcore.Origin(b"https", b"example.com", 443)
+
+        def is_available(self) -> bool:
+            return True  # Available with capacity
+
+        def is_idle(self) -> bool:
+            return False  # Not idle (active HTTP/2 connection with streams)
+
+        def is_closed(self) -> bool:
+            return False
+
+        def has_expired(self) -> bool:
+            return False
+
+        def get_available_stream_capacity(self) -> int:
+            return 5  # HTTP/2 connection with available stream capacity
+
+        def close(self):
+            pass
+
+    class MockIdleConnection(httpcore.ConnectionInterface):
+        def __init__(self) -> None:
+            self._origin = httpcore.Origin(b"https", b"example.com", 443)
+
+        def is_available(self) -> bool:
+            return True
+
+        def is_idle(self) -> bool:
+            return True  # This is an idle connection
+
+        def is_closed(self) -> bool:
+            return False
+
+        def has_expired(self) -> bool:
+            return False
+
+        def get_available_stream_capacity(self) -> int:
+            return 1
+
+    network_backend = httpcore.MockBackend([])
+
+    with httpcore.ConnectionPool(
+        max_keepalive_connections=0,  # Force keepalive limit to 0
+        http2=True,
+        network_backend=network_backend,
+    ) as pool:
+        # Create one idle connection and one active HTTP/2 connection
+        idle_conn = MockIdleConnection()
+        active_http2_conn = MockActiveHTTP2Connection()
+
+        with pool._optional_thread_lock:
+            pool._connections = [idle_conn, active_http2_conn]
+
+            # This should close idle connections but preserve the active HTTP/2 connection
+            closing_conns = pool._assign_requests_to_connections()
+
+            # Verify idle connection is marked for closing
+            assert len(closing_conns) == 1
+            assert idle_conn in closing_conns
+
+            # Verify the active HTTP/2 connection is preserved (line 421)
+            assert len(pool._connections) == 1
+            assert active_http2_conn in pool._connections