diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 3bdab1d231a..017e63ed1b8 100644 --- a/tests/unit/conftest.py +++ b/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() diff --git a/tests/unit/test_socket_policy.py b/tests/unit/test_socket_policy.py new file mode 100644 index 00000000000..f93794d1ba8 --- /dev/null +++ b/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()