refactor(transport): Simplify keepalive client factories

Author: CyMule

_test_unstructured_client/unit/test_split_collection_urls.py

@@ -6,6 +6,7 @@
 from pypdf import PdfWriter
 
 from unstructured_client import UnstructuredClient
+from unstructured_client._hooks.custom import split_pdf_hook
 from unstructured_client._hooks.custom.split_pdf_hook import _get_collection_url
 from unstructured_client.models import operations, shared
 from unstructured_client.utils import BackoffStrategy, RetryConfig
@@ -60,9 +61,9 @@ def respond(request):
     )
     # Split chunks create their own async HTTP client; keep every request offline.
     monkeypatch.setattr(
-        httpx,
-        "AsyncClient",
-        lambda **kwargs: original_async_client(**dict(kwargs, transport=transport, mounts={})),
+        split_pdf_hook,
+        "create_async_http_client",
+        lambda **kwargs: original_async_client(**kwargs, transport=transport),
     )
     return sdk, sent, sync_client, async_client
 

_test_unstructured_client/unit/test_split_pdf_hook.py

@@ -625,14 +625,13 @@ def test_unit_resolve_tls_config_client_cert_and_key_split(
 async def test_unit_run_tasks_forwards_tls_config_to_httpx_async_transport(
     monkeypatch: pytest.MonkeyPatch, tmp_path
 ):
-    """Split requests apply TLS and pool settings to the transport."""
+    """Split requests apply TLS settings to the transport."""
     ca_bundle = tmp_path / "ca.pem"
     ca_bundle.write_text("-----BEGIN CERTIFICATE-----\nfake\n-----END CERTIFICATE-----\n")
     monkeypatch.setenv("SSL_CERT_FILE", str(ca_bundle))
     monkeypatch.setenv("UNSTRUCTURED_CLIENT_TLS_CLIENT_CERT", "/etc/ssl/client.pem")
 
     monkeypatch.setenv("NO_PROXY", "*")
-    monkeypatch.setenv("UNSTRUCTURED_CLIENT_MAX_CONNECTIONS", "17")
     captured: dict = {}
 
     class _MockAsyncClient:
@@ -660,7 +659,6 @@ async def __aexit__(self, *exc):
     assert captured["transport"] is transport.return_value
     assert kwargs["verify"] == str(ca_bundle)
     assert kwargs["cert"] == "/etc/ssl/client.pem"
-    assert kwargs["limits"].max_connections == 17
     assert kwargs["socket_options"]
 
 

_test_unstructured_client/unit/test_tcp_keepalive.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import socket
-import sys
+from contextlib import contextmanager
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from threading import Thread
 from unittest.mock import patch
@@ -17,108 +17,127 @@
 
 @pytest.fixture(autouse=True)
 def clean_env(monkeypatch):
-    if hasattr(transport, "_log_keepalive_unavailable"):
-        transport._log_keepalive_unavailable.cache_clear()
     for suffix in ("", "_IDLE", "_INTERVAL", "_COUNT"):
         monkeypatch.delenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE" + suffix, raising=False)
     monkeypatch.setenv("NO_PROXY", "*")
 
 
-@pytest.mark.parametrize("platform", ["linux", "macos", "macos_old", "windows"])
-def test_unit_platform_options(monkeypatch, platform):
-    for name in ("TCP_KEEPIDLE", "TCP_KEEPALIVE", "TCP_KEEPINTVL", "TCP_KEEPCNT"):
-        monkeypatch.delattr(socket, name, raising=False)
-    expected = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
-    if platform != "windows":
-        monkeypatch.setattr(socket, "TCP_KEEPIDLE" if platform == "linux" else "TCP_KEEPALIVE", 100, raising=False)
-        expected.append((socket.IPPROTO_TCP, 100, 60))
-        if platform != "macos_old":
-            monkeypatch.setattr(socket, "TCP_KEEPINTVL", 101, raising=False)
-            monkeypatch.setattr(socket, "TCP_KEEPCNT", 102, raising=False)
-            expected += [(socket.IPPROTO_TCP, 101, 30), (socket.IPPROTO_TCP, 102, 5)]
-    assert transport.resolve_tcp_keepalive().socket_options() == expected
+DEFAULTS = transport.TCPKeepAlive(enabled=True)
+BASE_OPTIONS = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
 
 
-@pytest.mark.parametrize("value", ["0", "false", "False", " FALSE "])
-def test_unit_disabled(monkeypatch, value):
-    monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE", value)
-    assert transport.resolve_tcp_keepalive().socket_options() == []
+def _close(client):
+    if isinstance(client, httpx.AsyncClient):
+        asyncio.run(client.aclose())
+    else:
+        client.close()
 
 
-@pytest.mark.parametrize("suffix,default", [("IDLE", 60), ("INTERVAL", 30), ("COUNT", 5)])
-@pytest.mark.parametrize("value", ["0", "-1", "bad", "1.5", ""])
-def test_unit_invalid_value(monkeypatch, caplog, suffix, default, value):
-    monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix, value)
-    assert getattr(transport.resolve_tcp_keepalive(), suffix.lower()) == default
-    assert "using default" in caplog.text
+@contextmanager
+def local_server(handler):
+    server = ThreadingHTTPServer(("127.0.0.1", 0), handler)
+    thread = Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        yield server.server_port
+    finally:
+        server.shutdown()
+        server.server_close()
+        thread.join()
+
+
+@pytest.mark.parametrize("attrs_to_set,expected_options", [
+    ({"TCP_KEEPIDLE": 100, "TCP_KEEPINTVL": 101, "TCP_KEEPCNT": 102},
+     [(socket.IPPROTO_TCP, 100, DEFAULTS.idle),
+      (socket.IPPROTO_TCP, 101, DEFAULTS.interval),
+      (socket.IPPROTO_TCP, 102, DEFAULTS.count)]),
+    ({"TCP_KEEPALIVE": 100, "TCP_KEEPINTVL": 101, "TCP_KEEPCNT": 102},
+     [(socket.IPPROTO_TCP, 100, DEFAULTS.idle),
+      (socket.IPPROTO_TCP, 101, DEFAULTS.interval),
+      (socket.IPPROTO_TCP, 102, DEFAULTS.count)]),
+    ({"TCP_KEEPALIVE": 100}, [(socket.IPPROTO_TCP, 100, DEFAULTS.idle)]),
+    ({}, []),
+], ids=["linux", "macos", "macos_old", "windows"])
+def test_unit_platform_options(monkeypatch, attrs_to_set, expected_options):
+    for name in ("TCP_KEEPIDLE", "TCP_KEEPALIVE", "TCP_KEEPINTVL", "TCP_KEEPCNT"):
+        monkeypatch.delattr(socket, name, raising=False)
+    for name, value in attrs_to_set.items():
+        monkeypatch.setattr(socket, name, value, raising=False)
+    assert transport.resolve_tcp_keepalive().socket_options() == BASE_OPTIONS + expected_options
 
 
-@pytest.mark.parametrize("suffix,maximum,default", [
-    ("IDLE", 32767, 60), ("INTERVAL", 32767, 30), ("COUNT", 127, 5),
+@pytest.mark.parametrize("value", ["0", "false", "False", " FALSE "])
+def test_unit_disabled(monkeypatch, caplog, value):
+    monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE", value)
+    for suffix in ("IDLE", "INTERVAL", "COUNT"):
+        monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix, "invalid")
+    settings = transport.resolve_tcp_keepalive()
+    assert settings.socket_options() == []
+    assert settings.describe() == "off"
+    assert not caplog.records
+
+
+@pytest.mark.parametrize("suffix,value,expected", [
+    (suffix, value, expected)
+    for suffix, maximum in (
+        ("IDLE", transport.MAX_IDLE),
+        ("INTERVAL", transport.MAX_INTERVAL),
+        ("COUNT", transport.MAX_COUNT),
+    )
+    for value, expected in (
+        *[(invalid, getattr(DEFAULTS, suffix.lower()))
+          for invalid in ("0", "-1", "bad", "1.5", "", str(maximum + 1))],
+        ("1", 1),
+        (str(maximum), maximum),
+    )
 ])
-@pytest.mark.parametrize("above_maximum", [False, True])
-def test_unit_value_bounds(monkeypatch, caplog, suffix, maximum, default, above_maximum):
-    monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix,
-                       str(maximum + int(above_maximum)))
-    expected = default if above_maximum else maximum
+def test_unit_value_bounds(monkeypatch, caplog, suffix, value, expected):
+    monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix, value)
     assert getattr(transport.resolve_tcp_keepalive(), suffix.lower()) == expected
-    assert ("using default" in caplog.text) == above_maximum
+    assert ("using default" in caplog.text) == (value != str(expected))
 
 
 @pytest.mark.parametrize("asynchronous", [False, True])
 def test_unit_missing_proxy_helper_client_options(monkeypatch, asynchronous):
     monkeypatch.setattr(transport, "get_environment_proxies", None)
     if asynchronous:
-        options = {"verify": False, "cert": None, "limits": httpx.Limits(max_connections=3)}
+        options = {"verify": False, "cert": "client.pem", "limits": httpx.Limits(max_connections=3)}
         with patch.object(httpx, "AsyncClient") as constructor:
-            transport.create_async_http_client(timeout=42, transport_options=options)
+            transport.create_async_http_client(timeout=42, **options)
         constructor.assert_called_once_with(timeout=42, **options)
     else:
         with patch.object(httpx, "Client") as constructor:
             transport.create_http_client()
         constructor.assert_called_once_with()
 
 
-def test_unit_missing_proxy_helper_logs_once(monkeypatch, caplog):
-    monkeypatch.setattr(transport, "get_environment_proxies", None)
-    caplog.set_level("DEBUG", logger=transport.logger.name)
-    with transport.create_http_client():
-        pass
-    asyncio.run(transport.create_async_http_client().aclose())
-    messages = [record for record in caplog.records if "TCP keepalive is unavailable" in record.message]
-    assert len(messages) == 1
-    assert messages[0].levelname == "DEBUG"
-
-
 def test_unit_custom_values(monkeypatch):
     for suffix, value in (("IDLE", "5"), ("INTERVAL", "6"), ("COUNT", "7")):
         monkeypatch.setenv("UNSTRUCTURED_CLIENT_TCP_KEEPALIVE_" + suffix, value)
-    assert transport.resolve_tcp_keepalive() == transport.TCPKeepAlive(True, 5, 6, 7)
+    settings = transport.resolve_tcp_keepalive()
+    assert settings == transport.TCPKeepAlive(True, 5, 6, 7)
+    assert settings.describe() == "idle=5s interval=6s count=7"
 
 
 def test_unit_sdk_owned_clients():
     sdk = UnstructuredClient()
     options = transport.resolve_tcp_keepalive().socket_options()
     assert sdk.sdk_configuration.client._transport.base_transport._pool._socket_options == options
     assert sdk.sdk_configuration.async_client._transport._pool._socket_options == options
-    sdk.sdk_configuration.client.close()
-    asyncio.run(sdk.sdk_configuration.async_client.aclose())
+    _close(sdk.sdk_configuration.client)
+    _close(sdk.sdk_configuration.async_client)
 
 
 def test_unit_supplied_clients_untouched():
     sync = httpx.Client()
     asynchronous = httpx.AsyncClient()
     sync_transport = sync._transport
     async_transport = asynchronous._transport
-    sdk = UnstructuredClient(client=sync, async_client=asynchronous)
-    assert sdk.sdk_configuration.client is sync
-    assert sdk.sdk_configuration.async_client is asynchronous
+    UnstructuredClient(client=sync, async_client=asynchronous)
     assert sync._transport.base_transport is sync_transport
     assert asynchronous._transport is async_transport
-    assert sync_transport._pool._socket_options is None
-    assert asynchronous._transport._pool._socket_options is None
-    sync.close()
-    asyncio.run(asynchronous.aclose())
+    _close(sync)
+    _close(asynchronous)
 
 
 @pytest.mark.parametrize("missing_proxy_helper", [False, True])
@@ -132,21 +151,15 @@ def test_unit_environment_proxy_routes(monkeypatch, asynchronous, missing_proxy_
     proxied = client._transport_for_url(httpx.URL("https://api.example"))
     assert isinstance(proxied._pool, httpcore.AsyncHTTPProxy if asynchronous else httpcore.HTTPProxy)
     assert proxied._pool._proxy_url.host == b"127.0.0.1"
+    assert proxied._pool._socket_options is None
     if missing_proxy_helper:
         assert client._transport._pool._socket_options is None
-        assert proxied._pool._socket_options is None
     assert client._transport_for_url(httpx.URL("https://excluded.example")) is client._transport
-    if asynchronous:
-        asyncio.run(client.aclose())
-    else:
-        client.close()
+    _close(client)
 
 
 @pytest.mark.parametrize("asynchronous", [False, True])
 def test_unit_real_socket_keepalive(asynchronous):
-    if not hasattr(socket, "SO_KEEPALIVE"):
-        pytest.skip("SO_KEEPALIVE unavailable")
-
     class Handler(BaseHTTPRequestHandler):
         protocol_version = "HTTP/1.1"
 
@@ -156,39 +169,26 @@ def do_GET(self):
             self.end_headers()
             self.wfile.write(b"ok")
 
-    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
-    thread = Thread(target=server.serve_forever, daemon=True)
-    thread.start()
-    url = f"http://127.0.0.1:{server.server_port}"
-
     def check(response):
         stream = response.extensions["network_stream"]
         sock = stream.get_extra_info("socket")
-        enabled = sock.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE)
-        # Darwin reports the option bit mask rather than a normalized boolean.
-        if sys.platform == "darwin":
-            assert enabled != 0
-        else:
-            assert enabled == 1
+        assert sock.getsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE)
 
     async def request():
         async with transport.create_async_http_client() as client:
             async with client.stream("GET", url) as response:
                 check(response)
                 await response.aread()
 
-    try:
+    with local_server(Handler) as port:
+        url = f"http://127.0.0.1:{port}"
         if asynchronous:
             asyncio.run(request())
         else:
             with transport.create_http_client() as client:
                 with client.stream("GET", url) as response:
                     check(response)
                     response.read()
-    finally:
-        server.shutdown()
-        server.server_close()
-        thread.join()
 
 
 @pytest.mark.parametrize("asynchronous", [False, True])
@@ -200,25 +200,17 @@ def do_CONNECT(self):
             destinations.append(self.path)
             self.send_error(502)
 
-    server = ThreadingHTTPServer(("127.0.0.1", 0), ProxyHandler)
-    thread = Thread(target=server.serve_forever, daemon=True)
-    thread.start()
-    monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{server.server_port}")
-    monkeypatch.setenv("NO_PROXY", "")
-
     async def request():
         async with transport.create_async_http_client() as client:
             await client.get("https://proxy-target.invalid/")
 
-    try:
+    with local_server(ProxyHandler) as port:
+        monkeypatch.setenv("HTTPS_PROXY", f"http://127.0.0.1:{port}")
+        monkeypatch.setenv("NO_PROXY", "")
         with pytest.raises(httpx.ProxyError, match="502"):
             if asynchronous:
                 asyncio.run(request())
             else:
                 with transport.create_http_client() as client:
                     client.get("https://proxy-target.invalid/")
         assert destinations == ["proxy-target.invalid:443"]
-    finally:
-        server.shutdown()
-        server.server_close()
-        thread.join()