Merge pull request #42113 from BerriAI/litellm_unit_socket_block_at_import

test(unit): block external sockets at import time and add a socket policy regression test

tests/unit/conftest.py

@@ -1,9 +1,11 @@
-from collections.abc import Iterator
+import os
 from typing import Final
 
 import pytest
 from pytest_socket import enable_socket, socket_allow_hosts
 
+os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+
 LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1"]
 
 
@@ -11,13 +13,13 @@ def _allow_loopback_only() -> None:
     socket_allow_hosts(LOOPBACK_HOSTS, allow_unix_socket=True)
 
 
-@pytest.fixture(autouse=True, scope="session")
-def block_external_sockets() -> Iterator[None]:
-    _allow_loopback_only()
-    yield
-    enable_socket()
+_allow_loopback_only()
 
 
 @pytest.hookimpl(trylast=True)
 def pytest_runtest_setup() -> None:
     _allow_loopback_only()
+
+
+def pytest_sessionfinish() -> None:
+    enable_socket()

tests/unit/test_socket_policy.py

@@ -0,0 +1,17 @@
+import socket
+
+import pytest
+from pytest_socket import SocketConnectBlockedError
+
+
+def test_external_connect_is_refused_before_a_packet_leaves() -> None:
+    with pytest.raises(SocketConnectBlockedError):
+        socket.create_connection(("192.0.2.1", 9), timeout=1)
+
+
+def test_loopback_connect_is_allowed() -> None:
+    with socket.socket() as server:
+        server.bind(("127.0.0.1", 0))
+        server.listen()
+        with socket.create_connection(server.getsockname(), timeout=1) as client:
+            assert client.getpeername() == server.getsockname()