chore: import upstream snapshot with attribution

packages/python-sdk/tests/test_api_client_transport.py

@@ -0,0 +1,421 @@
+import asyncio
+import gc
+from concurrent.futures import ThreadPoolExecutor
+
+import httpx
+import pytest
+
+from e2b.api.client_async import AsyncEnvdTransportWithLogger, AsyncTransportWithLogger
+from e2b.api.client_async import get_api_client as get_async_api_client
+from e2b.api.client_async import get_envd_transport as get_async_envd_transport
+from e2b.api.client_async import get_transport as get_async_transport
+from e2b.api.client_sync import EnvdTransportWithLogger, TransportWithLogger
+from e2b.api.client_sync import get_api_client as get_sync_api_client
+from e2b.api.client_sync import get_envd_transport as get_sync_envd_transport
+from e2b.api.client_sync import get_transport as get_sync_transport
+from e2b.connection_config import ConnectionConfig
+
+
+def reset_sync_api_transports():
+    TransportWithLogger._thread_local.instances = {}
+
+
+def reset_sync_envd_transports():
+    EnvdTransportWithLogger._thread_local.instances = {}
+
+
+def run_in_worker_thread(fn):
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        return executor.submit(fn).result()
+
+
+def test_sync_api_client_proxy_uses_explicit_transport(test_api_key):
+    reset_sync_api_transports()
+    config = ConnectionConfig(
+        api_key=test_api_key,
+        proxy="http://127.0.0.1:9999",
+    )
+
+    api_client = get_sync_api_client(config)
+    httpx_client = api_client.get_httpx_client()
+
+    try:
+        assert "proxy" not in api_client._httpx_args
+        assert httpx_client._transport is get_sync_transport(config)
+        assert httpx_client._mounts == {}
+    finally:
+        httpx_client.close()
+        reset_sync_api_transports()
+
+
+def test_sync_get_transport_http2_opt_out_returns_distinct_instance(test_api_key):
+    reset_sync_api_transports()
+    config = ConnectionConfig(api_key=test_api_key)
+
+    try:
+        http2_transport = get_sync_transport(config)
+        http1_transport = get_sync_transport(config, http2=False)
+
+        assert http2_transport is not http1_transport
+        assert http2_transport._pool._http2 is True
+        assert http1_transport._pool._http2 is False
+        # Subsequent calls with the same http2 flag return the cached
+        # instance.
+        assert get_sync_transport(config) is http2_transport
+        assert get_sync_transport(config, http2=False) is http1_transport
+    finally:
+        reset_sync_api_transports()
+
+
+def test_sync_get_transport_keyed_by_proxy(test_api_key):
+    reset_sync_api_transports()
+    proxied_config = ConnectionConfig(
+        api_key=test_api_key,
+        proxy="http://127.0.0.1:9999",
+    )
+    direct_config = ConnectionConfig(api_key=test_api_key)
+    other_proxy_config = ConnectionConfig(
+        api_key=test_api_key,
+        proxy="http://127.0.0.1:8888",
+    )
+
+    try:
+        proxied_transport = get_sync_transport(proxied_config)
+        direct_transport = get_sync_transport(direct_config)
+        other_proxy_transport = get_sync_transport(other_proxy_config)
+
+        assert proxied_transport is not direct_transport
+        assert proxied_transport is not other_proxy_transport
+        assert direct_transport is not other_proxy_transport
+        # The same proxy still reuses the cached instance.
+        assert get_sync_transport(proxied_config) is proxied_transport
+        assert get_sync_transport(direct_config) is direct_transport
+    finally:
+        reset_sync_api_transports()
+
+
+def test_sync_api_client_applies_request_timeout(test_api_key):
+    reset_sync_api_transports()
+    config = ConnectionConfig(api_key=test_api_key, request_timeout=1.5)
+
+    api_client = get_sync_api_client(config)
+    httpx_client = api_client.get_httpx_client()
+
+    try:
+        assert httpx_client.timeout == httpx.Timeout(1.5)
+    finally:
+        httpx_client.close()
+        reset_sync_api_transports()
+
+
+def test_sync_api_client_request_timeout_zero_disables_timeout(test_api_key):
+    reset_sync_api_transports()
+    config = ConnectionConfig(api_key=test_api_key, request_timeout=0)
+
+    api_client = get_sync_api_client(config)
+    httpx_client = api_client.get_httpx_client()
+
+    try:
+        assert httpx_client.timeout == httpx.Timeout(None)
+    finally:
+        httpx_client.close()
+        reset_sync_api_transports()
+
+
+def test_sync_envd_transport_uses_separate_cache(test_api_key):
+    reset_sync_api_transports()
+    reset_sync_envd_transports()
+    config = ConnectionConfig(api_key=test_api_key)
+
+    try:
+        api_transport = get_sync_transport(config)
+        envd_transport = get_sync_envd_transport(config)
+
+        assert api_transport is not envd_transport
+        assert get_sync_transport(config) is api_transport
+        assert get_sync_envd_transport(config) is envd_transport
+        assert envd_transport._pool._http2 is True
+    finally:
+        reset_sync_api_transports()
+        reset_sync_envd_transports()
+
+
+def test_sync_api_transport_cache_reuses_within_thread_and_isolates_across_threads(
+    test_api_key,
+):
+    reset_sync_api_transports()
+    config = ConnectionConfig(api_key=test_api_key)
+
+    try:
+        main_transport = get_sync_transport(config)
+        same_thread_transport = get_sync_transport(config)
+        worker_thread_transport = run_in_worker_thread(
+            lambda: get_sync_transport(config)
+        )
+
+        assert same_thread_transport is main_transport
+        assert worker_thread_transport is not main_transport
+    finally:
+        reset_sync_api_transports()
+
+
+def test_sync_api_client_cache_reuses_within_thread_and_isolates_across_threads(
+    test_api_key,
+):
+    reset_sync_api_transports()
+    config = ConnectionConfig(api_key=test_api_key)
+    api_client = get_sync_api_client(config)
+
+    def get_worker_client_ids():
+        httpx_client = api_client.get_httpx_client()
+        try:
+            return (
+                id(httpx_client),
+                id(api_client.get_httpx_client()),
+                id(httpx_client._transport),
+                id(get_sync_transport(config)),
+            )
+        finally:
+            httpx_client.close()
+
+    try:
+        main_client = api_client.get_httpx_client()
+        (
+            worker_client_id,
+            worker_cached_client_id,
+            worker_transport_id,
+            worker_cached_transport_id,
+        ) = run_in_worker_thread(get_worker_client_ids)
+
+        assert api_client.get_httpx_client() is main_client
+        assert worker_client_id == worker_cached_client_id
+        assert worker_transport_id == worker_cached_transport_id
+        assert worker_client_id != id(main_client)
+        assert worker_transport_id != id(main_client._transport)
+    finally:
+        main_client.close()
+        reset_sync_api_transports()
+
+
+def test_sync_envd_transport_cache_is_thread_local(test_api_key):
+    reset_sync_envd_transports()
+    config = ConnectionConfig(api_key=test_api_key)
+
+    try:
+        main_transport = get_sync_envd_transport(config)
+        thread_transport = run_in_worker_thread(lambda: get_sync_envd_transport(config))
+
+        assert main_transport is get_sync_envd_transport(config)
+        assert thread_transport is not main_transport
+        assert main_transport._pool._http2 is True
+        assert thread_transport._pool._http2 is True
+    finally:
+        reset_sync_envd_transports()
+
+
+@pytest.mark.asyncio
+async def test_async_api_client_proxy_uses_explicit_transport(test_api_key):
+    AsyncTransportWithLogger._instances.clear()
+    config = ConnectionConfig(
+        api_key=test_api_key,
+        proxy="http://127.0.0.1:9999",
+    )
+
+    api_client = get_async_api_client(config)
+    httpx_client = api_client.get_async_httpx_client()
+    transport = AsyncTransportWithLogger._instances[asyncio.get_running_loop()][
+        (True, "http://127.0.0.1:9999")
+    ]
+
+    try:
+        assert "proxy" not in api_client._httpx_args
+        assert httpx_client._transport is transport
+        assert httpx_client._mounts == {}
+    finally:
+        await httpx_client.aclose()
+        AsyncTransportWithLogger._instances.clear()
+
+
+@pytest.mark.asyncio
+async def test_async_get_transport_http2_opt_out_returns_distinct_instance(
+    test_api_key,
+):
+    AsyncTransportWithLogger._instances.clear()
+    config = ConnectionConfig(api_key=test_api_key)
+
+    try:
+        http2_transport = get_async_transport(config)
+        http1_transport = get_async_transport(config, http2=False)
+
+        assert http2_transport is not http1_transport
+        assert http2_transport._pool._http2 is True
+        assert http1_transport._pool._http2 is False
+        # Subsequent calls with the same http2 flag return the cached
+        # instance.
+        assert get_async_transport(config) is http2_transport
+        assert get_async_transport(config, http2=False) is http1_transport
+    finally:
+        AsyncTransportWithLogger._instances.clear()
+
+
+@pytest.mark.asyncio
+async def test_async_get_transport_keyed_by_proxy(test_api_key):
+    AsyncTransportWithLogger._instances.clear()
+    proxied_config = ConnectionConfig(
+        api_key=test_api_key,
+        proxy="http://127.0.0.1:9999",
+    )
+    direct_config = ConnectionConfig(api_key=test_api_key)
+
+    try:
+        proxied_transport = get_async_transport(proxied_config)
+        direct_transport = get_async_transport(direct_config)
+
+        assert proxied_transport is not direct_transport
+        # The same proxy still reuses the cached instance.
+        assert get_async_transport(proxied_config) is proxied_transport
+        assert get_async_transport(direct_config) is direct_transport
+    finally:
+        AsyncTransportWithLogger._instances.clear()
+
+
+@pytest.mark.asyncio
+async def test_async_api_client_applies_request_timeout(test_api_key):
+    AsyncTransportWithLogger._instances.clear()
+    config = ConnectionConfig(api_key=test_api_key, request_timeout=1.5)
+
+    api_client = get_async_api_client(config)
+    httpx_client = api_client.get_async_httpx_client()
+
+    try:
+        assert httpx_client.timeout == httpx.Timeout(1.5)
+    finally:
+        await httpx_client.aclose()
+        AsyncTransportWithLogger._instances.clear()
+
+
+@pytest.mark.asyncio
+async def test_async_api_client_cache_reuses_within_loop_and_isolates_across_loops(
+    test_api_key,
+):
+    AsyncTransportWithLogger._instances.clear()
+    config = ConnectionConfig(api_key=test_api_key)
+    api_client = get_async_api_client(config)
+
+    async def get_client_ids():
+        httpx_client = api_client.get_async_httpx_client()
+        try:
+            return (
+                id(httpx_client),
+                id(api_client.get_async_httpx_client()),
+                id(httpx_client._transport),
+                id(get_async_transport(config)),
+            )
+        finally:
+            await httpx_client.aclose()
+
+    try:
+        main_client = api_client.get_async_httpx_client()
+        (
+            worker_client_id,
+            worker_cached_client_id,
+            worker_transport_id,
+            worker_cached_transport_id,
+        ) = await asyncio.get_running_loop().run_in_executor(
+            None,
+            lambda: asyncio.run(get_client_ids()),
+        )
+
+        assert api_client.get_async_httpx_client() is main_client
+        assert worker_client_id == worker_cached_client_id
+        assert worker_transport_id == worker_cached_transport_id
+        assert worker_client_id != id(main_client)
+        assert worker_transport_id != id(main_client._transport)
+    finally:
+        await main_client.aclose()
+        AsyncTransportWithLogger._instances.clear()
+
+
+def test_async_transport_not_reused_across_sequential_loops(test_api_key):
+    AsyncTransportWithLogger._instances.clear()
+    config = ConnectionConfig(api_key=test_api_key)
+
+    async def get_transport():
+        return get_async_transport(config)
+
+    try:
+        loop_a = asyncio.new_event_loop()
+        try:
+            transport_a = loop_a.run_until_complete(get_transport())
+        finally:
+            loop_a.close()
+        del loop_a
+        gc.collect()
+
+        # The cache entry dies with the loop, so a later loop can never
+        # inherit a transport bound to a closed loop, even when CPython
+        # reuses the dead loop's object id.
+        assert len(AsyncTransportWithLogger._instances) == 0
+
+        loop_b = asyncio.new_event_loop()
+        try:
+            transport_b = loop_b.run_until_complete(get_transport())
+        finally:
+            loop_b.close()
+
+        assert transport_b is not transport_a
+    finally:
+        AsyncTransportWithLogger._instances.clear()
+
+
+def test_async_api_client_not_reused_across_sequential_loops(test_api_key):
+    AsyncTransportWithLogger._instances.clear()
+    config = ConnectionConfig(api_key=test_api_key)
+    api_client = get_async_api_client(config)
+
+    async def get_client():
+        client = api_client.get_async_httpx_client()
+        assert api_client.get_async_httpx_client() is client
+        return client
+
+    try:
+        loop_a = asyncio.new_event_loop()
+        try:
+            client_a = loop_a.run_until_complete(get_client())
+            loop_a.run_until_complete(client_a.aclose())
+        finally:
+            loop_a.close()
+        del loop_a
+        gc.collect()
+
+        assert len(api_client._async_clients) == 0
+
+        loop_b = asyncio.new_event_loop()
+        try:
+            client_b = loop_b.run_until_complete(get_client())
+            loop_b.run_until_complete(client_b.aclose())
+        finally:
+            loop_b.close()
+
+        assert client_b is not client_a
+    finally:
+        AsyncTransportWithLogger._instances.clear()
+
+
+@pytest.mark.asyncio
+async def test_async_envd_transport_uses_separate_cache(test_api_key):
+    AsyncTransportWithLogger._instances.clear()
+    AsyncEnvdTransportWithLogger._instances.clear()
+    config = ConnectionConfig(api_key=test_api_key)
+
+    try:
+        api_transport = get_async_transport(config)
+        envd_transport = get_async_envd_transport(config)
+
+        assert api_transport is not envd_transport
+        assert get_async_transport(config) is api_transport
+        assert get_async_envd_transport(config) is envd_transport
+        assert envd_transport._pool._http2 is True
+    finally:
+        AsyncTransportWithLogger._instances.clear()
+        AsyncEnvdTransportWithLogger._instances.clear()