test(lens): run connection reuse coverage as an integration test

.github/workflows/test-unit.yml

@@ -483,6 +483,7 @@ jobs:
               tests/unit/sandbox
               tests/unit/skills/test_skills_main.py
               tests/unit/tracing
+              tests/proxy_behavior/lens/test_connection.py
             workers: 2
             reruns: 0
             timeout-minutes: 20

tests/proxy_behavior/lens/test_connection.py

@@ -0,0 +1,46 @@
+import asyncio
+from typing import Final
+
+import pytest
+
+from litellm.tracing.remote import LensConnection
+
+
+@pytest.mark.asyncio
+async def test_control_requests_reuse_connections_without_retaining_another_service_credential() -> None:
+    requests: Final[asyncio.Queue[tuple[str, bytes]]] = asyncio.Queue()
+
+    async def serve(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        try:
+            while True:
+                headers: Final = await reader.readuntil(b"\r\n\r\n")
+                requests.put_nowait((str(writer.get_extra_info("peername")), headers))
+                writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}")
+                await writer.drain()
+        except asyncio.IncompleteReadError:
+            pass
+        finally:
+            writer.close()
+            await writer.wait_closed()
+
+    async with await asyncio.start_server(serve, "127.0.0.1", 0) as server:
+        port: Final = server.sockets[0].getsockname()[1]
+        first: Final = LensConnection(f"http://127.0.0.1:{port}/one", "first-service-token")
+        second: Final = LensConnection(f"http://127.0.0.1:{port}/two", "second-service-token")
+        try:
+            for connection in (first, second):
+                response: Final = await connection.control_client().get(
+                    connection.endpoint("/internal/status"), headers=connection.headers
+                )
+                assert response.json() == {}
+            first_peer, first_request = await asyncio.wait_for(requests.get(), 2)
+            second_peer, second_request = await asyncio.wait_for(requests.get(), 2)
+            assert first_peer == second_peer
+            assert b"GET /one/internal/status " in first_request
+            assert b"GET /two/internal/status " in second_request
+            assert b"Bearer first-service-token" in first_request
+            assert b"Bearer second-service-token" not in first_request
+            assert b"Bearer second-service-token" in second_request
+            assert b"Bearer first-service-token" not in second_request
+        finally:
+            await second.control_client().aclose()

tests/unit/tracing/test_remote.py

@@ -185,43 +185,3 @@ async def test_request_records_use_the_internal_service_endpoint() -> None:
     request: Final = requests.get_nowait()
     assert request.url.path == "/internal/spend"
     assert json.loads(request.content) == [{"request_id": "r"}]
-
-
-@pytest.mark.asyncio
-async def test_control_requests_reuse_connections_without_retaining_another_service_credential() -> None:
-    requests: Final[asyncio.Queue[tuple[str, bytes]]] = asyncio.Queue()
-
-    async def serve(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
-        try:
-            while True:
-                headers: Final = await reader.readuntil(b"\r\n\r\n")
-                requests.put_nowait((str(writer.get_extra_info("peername")), headers))
-                writer.write(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\n{}")
-                await writer.drain()
-        except asyncio.IncompleteReadError:
-            pass
-        finally:
-            writer.close()
-            await writer.wait_closed()
-
-    async with await asyncio.start_server(serve, "127.0.0.1", 0) as server:
-        port: Final = server.sockets[0].getsockname()[1]
-        first: Final = LensConnection(f"http://127.0.0.1:{port}/one", "first-service-token")
-        second: Final = LensConnection(f"http://127.0.0.1:{port}/two", "second-service-token")
-        try:
-            for connection in (first, second):
-                response: Final = await connection.control_client().get(
-                    connection.endpoint("/internal/status"), headers=connection.headers
-                )
-                assert response.json() == {}
-            first_peer, first_request = await asyncio.wait_for(requests.get(), 2)
-            second_peer, second_request = await asyncio.wait_for(requests.get(), 2)
-            assert first_peer == second_peer
-            assert b"GET /one/internal/status " in first_request
-            assert b"GET /two/internal/status " in second_request
-            assert b"Bearer first-service-token" in first_request
-            assert b"Bearer second-service-token" not in first_request
-            assert b"Bearer second-service-token" in second_request
-            assert b"Bearer first-service-token" not in second_request
-        finally:
-            await second.control_client().aclose()