mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge remote-tracking branch 'origin/main' into litellm_agent_access_groups
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
commit
85da9cb588
2 changed files with 25 additions and 6 deletions
|
|
@ -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()
|
||||
|
|
|
|||
17
tests/unit/test_socket_policy.py
Normal file
17
tests/unit/test_socket_policy.py
Normal file
|
|
@ -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()
|
||||
Loading…
Add table
Reference in a new issue