scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 267

test(unit): keep unix socket paths in unit tests under the macOS AF_UNIX limit (#45040)

Browse source 
* test(proxy): keep the spend sidecar socket path under the macOS AF_UNIX limit

* test(unit): share the short socket fixtures across unit tests

Move socket_dir and socket_path to tests/unit/conftest.py so the datadog
redis test's absent socket fails for the right reason on a Mac, rename
the pgbouncer helper's directory argument, and record the rule in
tests/unit/AGENTS.md

* test(unit): fall back to /tmp when TMPDIR leaves no room for a socket path

socket_dir now uses /tmp when TMPDIR plus the longest socket suffix
would pass macOS's sun_path limit, and its teardown ignores errors so a
failed test reports only its own failure

* test(unit): keep socket fixtures inside tmp_path behind a short link

socket_dir now links a short temp directory to tmp_path, so artifacts keep
pytest's retention and a killed run leaks only an empty link. The fallback
to /tmp counts the temp directory in bytes, and socket_tempdir points
tempfile at it for the pooler tests whose socket directory production picks

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-28 20:23:22 -07:00 • committed by GitHub
parent 56e60d3b7d
commit a2ef18b257
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 142 additions and 101 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
tests/unit/AGENTS.md
View file

@ -8,6 +8,10 @@ the same tests against Google's live spec
Unit tests must not depend on the current time or execution speed. Use fixed timestamps for fixtures
and controlled clocks for time-dependent behavior. Incidental clock reads inside production code are allowed
A unix socket goes under the `socket_dir` or `socket_path` fixture, never `tmp_path`, whose paths run past macOS's
104-byte AF_UNIX limit and fail only on a Mac. Code that opens a socket in its own `tempfile` directory runs under
`socket_tempdir`
## What good looks like
```python

34
tests/unit/conftest.py
View file

@ -2,6 +2,8 @@ import asyncio
import base64
import importlib
import os
import shutil
import tempfile
from collections.abc import Callable, Coroutine, Iterator
from dataclasses import dataclass, field
from pathlib import Path
@ -217,6 +219,38 @@ def _flush_client_caches() -> None:
_reset_aws_auth_caches()
_MACOS_LONGEST_SOCKET_PATH: Final = 103
_LONGEST_SOCKET_SUFFIX: Final = len("/sock12345678/d/litellm-pgbouncer-12345678/.s.PGSQL.65535")
def _socket_root() -> str:
default: Final = tempfile.gettempdir()
fits: Final = len(os.fsencode(default)) + _LONGEST_SOCKET_SUFFIX <= _MACOS_LONGEST_SOCKET_PATH
return default if fits else "/tmp"
@pytest.fixture
def socket_dir(tmp_path: Path) -> Iterator[Path]:
"""``tmp_path`` behind a short link, so unix socket paths fit the sun_path limit that it exceeds on macOS."""
link_dir: Final = Path(tempfile.mkdtemp(prefix="sock", dir=_socket_root()))
link: Final = link_dir / "d"
link.symlink_to(tmp_path, target_is_directory=True)
yield link
shutil.rmtree(link_dir, ignore_errors=True)
@pytest.fixture
def socket_path(socket_dir: Path) -> Path:
return socket_dir / "spend.sock"
@pytest.fixture
def socket_tempdir(socket_dir: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
"""Point ``tempfile``'s default directory at ``socket_dir``, for code that picks its own socket directory."""
monkeypatch.setattr(tempfile, "tempdir", str(socket_dir))
return socket_dir
@pytest.fixture(autouse=True, scope="session")
def bundled_tiktoken_cache() -> None:
importlib.import_module("litellm.litellm_core_utils.default_encoding")

12
tests/unit/gateway/test_launch.py
View file

@ -140,7 +140,7 @@ class TestPoolDatabaseUrl:
assert "DATABASE_URL" in outcome.reason
def test_token_auth_hands_the_workers_the_pool_user_not_the_token(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, socket_tempdir: Path
):
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
@ -149,7 +149,7 @@ class TestPoolDatabaseUrl:
with _minted_iam_token("MINTED_TOKEN"):
outcome: Final = pool_database_url(
DatabaseURLSettings.from_env(),
PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path))),
PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(socket_tempdir))),
environ,
)
assert isinstance(outcome, str), outcome
@ -160,12 +160,12 @@ class TestPoolDatabaseUrl:
class TestMain:
def test_workers_inherit_the_loopback_url_the_supervisor_installed(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, socket_tempdir: Path
):
port: Final = _free_port()
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port))
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path)))
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(socket_tempdir)))
monkeypatch.setenv("KEEPALIVE_TIMEOUT", "75")
served: Final[list[tuple[str, ...]]] = []
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
@ -183,7 +183,7 @@ class TestMain:
assert _query(os.environ["DATABASE_URL"])["pgbouncer"] == "true"
def test_iam_workers_keep_the_loopback_url_instead_of_minting_their_own(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, socket_tempdir: Path
):
port: Final = _free_port()
monkeypatch.delenv("DATABASE_PASSWORD")
@ -191,7 +191,7 @@ class TestMain:
monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port))
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path)))
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(socket_tempdir)))
served: Final[list[tuple[str, ...]]] = []
with _minted_iam_token("SUPERVISOR_TOKEN"):
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))

4
tests/unit/integrations/datadog/test_datadog.py
View file

@ -905,10 +905,10 @@ async def test_a_failed_request_is_delivered_as_an_error_log_that_keeps_the_erro
@pytest.mark.asyncio
async def test_a_failing_redis_cache_is_delivered_to_datadog_as_redis_warnings(
delivery: tuple[DataDogLogger, respx.Route], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
delivery: tuple[DataDogLogger, respx.Route], monkeypatch: pytest.MonkeyPatch, socket_dir: Path
) -> None:
datadog_logger, intake = delivery
absent_socket: Final = str(tmp_path / "absent.sock")
absent_socket: Final = str(socket_dir / "absent.sock")
redis_cache: Final = Cache(type="redis", url=f"unix://{absent_socket}")
monkeypatch.setattr(redis_cache.cache.service_logger_obj, "dd_logger", datadog_logger, raising=False)
monkeypatch.setattr(litellm, "service_callback", ["datadog"])

122
tests/unit/proxy/db/test_pgbouncer.py
View file

@ -277,7 +277,7 @@ def _free_port() -> int:
def _fake_pooler(
tmp_path: Path,
directory: Path,
port: int,
exit_immediately: bool = False,
port_file: Path | None = None,
@ -288,13 +288,13 @@ def _fake_pooler(
"""An executable that listens like PgBouncer: on the TCP port first, then on ``.s.PGSQL.<port>`` in the socket dir.
Port and socket dir come from the ini it is given, else from ``port`` and
``tmp_path``. With ``port_file`` each start reads the port from that file
``directory``. With ``port_file`` each start reads the port from that file
instead. ``bind_delay_seconds`` holds the bind back, like a slow start.
``--version`` prints ``version_banner``. With ``auth_log`` it appends the
``auth_file`` it reads at startup and on every SIGHUP, one line per read,
like PgBouncer loading its credentials.
"""
script: Final = tmp_path / "fake-pgbouncer"
script: Final = directory / "fake-pgbouncer"
script.write_text(
textwrap.dedent(
f"""\
@ -316,7 +316,7 @@ def _fake_pooler(
port = ini.getint("pgbouncer", "listen_port", fallback={port})
if not {port_file is None!r}:
port = int(pathlib.Path({str(port_file)!r}).read_text())
socket_dir = ini.get("pgbouncer", "unix_socket_dir", fallback={str(tmp_path)!r})
socket_dir = ini.get("pgbouncer", "unix_socket_dir", fallback={str(directory)!r})
socket_path = f"{{socket_dir}}/.s.PGSQL.{{port}}"
time.sleep({bind_delay_seconds!r})
listener = socket.socket()
@ -357,10 +357,10 @@ def _wait_until(condition: Callable[[], bool], timeout_seconds: float = 5.0) ->
class TestPgBouncerProcess:
def test_start_waits_for_the_listener_and_stop_ends_it(self, tmp_path: Path):
def test_start_waits_for_the_listener_and_stop_ends_it(self, socket_dir: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port)
argv=(str(_fake_pooler(socket_dir, port)),), port=port, socket_path=unix_socket_path(socket_dir, port)
)
assert pooler.start() is None
assert _listening(port)
@ -371,12 +371,12 @@ class TestPgBouncerProcess:
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def test_a_crashed_pooler_is_restarted_with_a_new_pid(self, tmp_path: Path):
def test_a_crashed_pooler_is_restarted_with_a_new_pid(self, socket_dir: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
argv=(str(_fake_pooler(socket_dir, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
@ -388,17 +388,17 @@ class TestPgBouncerProcess:
assert _wait_until(lambda: not _listening(port))
def test_a_failed_restart_is_retried_until_the_pooler_is_back(
self, tmp_path: Path, caplog: pytest.LogCaptureFixture
self, socket_dir: Path, caplog: pytest.LogCaptureFixture
):
port: Final = _free_port()
script: Final = _fake_pooler(tmp_path, port)
script: Final = _fake_pooler(socket_dir, port)
pooler: Final = PgBouncerProcess(
argv=(str(script),), port=port, socket_path=unix_socket_path(tmp_path, port), restart_delay_seconds=0.1
argv=(str(script),), port=port, socket_path=unix_socket_path(socket_dir, port), restart_delay_seconds=0.1
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
hidden: Final = script.rename(tmp_path / "hidden")
hidden: Final = script.rename(socket_dir / "hidden")
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: any("could not be restarted" in record.message for record in caplog.records))
@ -408,15 +408,17 @@ class TestPgBouncerProcess:
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_replacement_that_never_listens_is_replaced_again(self, tmp_path: Path, caplog: pytest.LogCaptureFixture):
def test_a_replacement_that_never_listens_is_replaced_again(
self, socket_dir: Path, caplog: pytest.LogCaptureFixture
):
port: Final = _free_port()
port_file: Final = tmp_path / "port"
port_file: Final = socket_dir / "port"
port_file.write_text(str(port))
script: Final = _fake_pooler(tmp_path, port, port_file=port_file)
script: Final = _fake_pooler(socket_dir, port, port_file=port_file)
pooler: Final = PgBouncerProcess(
argv=(str(script),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
restart_delay_seconds=0.1,
ready_timeout_seconds=2.0,
)
@ -435,12 +437,12 @@ class TestPgBouncerProcess:
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_stopping_during_the_restart_delay_leaves_no_pooler_behind(self, tmp_path: Path):
def test_stopping_during_the_restart_delay_leaves_no_pooler_behind(self, socket_dir: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
argv=(str(_fake_pooler(socket_dir, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
restart_delay_seconds=0.3,
)
assert pooler.start() is None
@ -453,12 +455,12 @@ class TestPgBouncerProcess:
assert not _listening(port)
assert pooler.pid == first_pid
def test_a_stopped_pooler_is_not_restarted(self, tmp_path: Path, caplog: pytest.LogCaptureFixture):
def test_a_stopped_pooler_is_not_restarted(self, socket_dir: Path, caplog: pytest.LogCaptureFixture):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
argv=(str(_fake_pooler(socket_dir, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
@ -468,31 +470,31 @@ class TestPgBouncerProcess:
assert not _listening(port)
assert caplog.records == []
def test_a_pooler_that_exits_during_startup_is_reported(self, tmp_path: Path):
def test_a_pooler_that_exits_during_startup_is_reported(self, socket_dir: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, exit_immediately=True)),),
argv=(str(_fake_pooler(socket_dir, port, exit_immediately=True)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
assert "status 3" in outcome.reason
def test_a_missing_binary_is_reported(self, tmp_path: Path):
def test_a_missing_binary_is_reported(self, socket_dir: Path):
outcome: Final = PgBouncerProcess(
argv=("/nonexistent/pgbouncer",), port=_free_port(), socket_path=tmp_path / "sock"
argv=("/nonexistent/pgbouncer",), port=_free_port(), socket_path=socket_dir / "sock"
).start()
assert isinstance(outcome, PgBouncerError)
assert "/nonexistent/pgbouncer" in outcome.reason
def test_a_port_owned_by_someone_else_is_refused_before_spawning(self, tmp_path: Path):
def test_a_port_owned_by_someone_else_is_refused_before_spawning(self, socket_dir: Path):
with socket.socket() as squatter:
squatter.bind(("127.0.0.1", 0))
squatter.listen()
port: Final = _bound_port(squatter)
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port)
argv=(str(_fake_pooler(socket_dir, port)),), port=port, socket_path=unix_socket_path(socket_dir, port)
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
@ -500,13 +502,13 @@ class TestPgBouncerProcess:
assert pooler.pid is None
def test_a_replacement_waits_until_a_squatter_leaves_the_port(
self, tmp_path: Path, caplog: pytest.LogCaptureFixture
self, socket_dir: Path, caplog: pytest.LogCaptureFixture
):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
argv=(str(_fake_pooler(socket_dir, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
restart_delay_seconds=0.5,
)
assert pooler.start() is None
@ -524,12 +526,12 @@ class TestPgBouncerProcess:
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_listener_that_grabs_the_port_after_the_spawn_is_not_taken_for_the_pooler(self, tmp_path: Path):
def test_a_listener_that_grabs_the_port_after_the_spawn_is_not_taken_for_the_pooler(self, socket_dir: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=0.5)),),
argv=(str(_fake_pooler(socket_dir, port, bind_delay_seconds=0.5)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
ready_timeout_seconds=3.0,
)
with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter:
@ -542,12 +544,12 @@ class TestPgBouncerProcess:
assert isinstance(outcome, PgBouncerError)
assert "exited with status 1" in outcome.reason
def test_a_port_served_by_a_stranger_while_the_pooler_is_still_starting_is_reported(self, tmp_path: Path):
def test_a_port_served_by_a_stranger_while_the_pooler_is_still_starting_is_reported(self, socket_dir: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=30.0)),),
argv=(str(_fake_pooler(socket_dir, port, bind_delay_seconds=30.0)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
ready_timeout_seconds=0.5,
)
with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter:
@ -564,12 +566,12 @@ class TestPgBouncerProcess:
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def test_a_pooler_that_never_listens_times_out(self, tmp_path: Path):
def test_a_pooler_that_never_listens_times_out(self, socket_dir: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, _free_port())),),
argv=(str(_fake_pooler(socket_dir, _free_port())),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
ready_timeout_seconds=0.5,
)
outcome: Final = pooler.start()
@ -704,10 +706,10 @@ def _runtime_dir_listening_on(port: int) -> Path:
class TestStartInContainerPgBouncer:
def test_returns_the_loopback_url_once_the_pooler_listens(self, tmp_path: Path):
def test_returns_the_loopback_url_once_the_pooler_listens(self, socket_tempdir: Path):
port: Final = _free_port()
auth_log: Final = tmp_path / "auth.log"
binary: Final = _fake_pooler(tmp_path, port, auth_log=auth_log)
auth_log: Final = socket_tempdir / "auth.log"
binary: Final = _fake_pooler(socket_tempdir, port, auth_log=auth_log)
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm?connection_limit=5")
assert isinstance(pooled, str), pooled
@ -723,9 +725,9 @@ class TestStartInContainerPgBouncer:
assert auth_log.read_text() == repr(f'"app" "pw"\n"litellm_pgbouncer" "{parsed.password}"\n') + "\n"
@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning")
def test_a_forked_worker_exiting_leaves_the_pooler_and_its_files_to_the_parent(self, tmp_path: Path):
def test_a_forked_worker_exiting_leaves_the_pooler_and_its_files_to_the_parent(self, socket_tempdir: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(socket_tempdir, port)))
exit_hooks: Final[list[Callable[[], None]]] = []
pooled: Final = start_in_container_pgbouncer(
settings, "postgresql://app:pw@db/litellm", register_exit_hook=exit_hooks.append
@ -766,10 +768,10 @@ class TestStartInContainerPgBouncer:
assert "IAM_TOKEN_DB_AUTH" in outcome.reason
assert not _listening(port)
def test_token_auth_mints_the_first_token_into_the_auth_file_before_the_pooler_starts(self, tmp_path: Path):
def test_token_auth_mints_the_first_token_into_the_auth_file_before_the_pooler_starts(self, socket_tempdir: Path):
port: Final = _free_port()
auth_log: Final = tmp_path / "auth.log"
binary: Final = _fake_pooler(tmp_path, port, auth_log=auth_log)
auth_log: Final = socket_tempdir / "auth.log"
binary: Final = _fake_pooler(socket_tempdir, port, auth_log=auth_log)
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
token: Final = _entra_jwt(datetime.now(tz=timezone.utc) + timedelta(hours=1))
pooled: Final = start_in_container_pgbouncer(
@ -798,25 +800,25 @@ class TestStartInContainerPgBouncer:
assert "no Azure credential" in outcome.reason
assert not _listening(port)
def test_a_renewed_token_is_written_and_picked_up_by_the_running_and_by_a_restarted_pooler(self, tmp_path: Path):
def test_a_renewed_token_is_written_and_picked_up_by_the_running_and_by_a_restarted_pooler(self, socket_dir: Path):
port: Final = _free_port()
auth_log: Final = tmp_path / "auth.log"
auth_log: Final = socket_dir / "auth.log"
plan: Final = plan_pgbouncer(
"postgresql://app@db/litellm", PgBouncerSettings(enabled=True, port=port), tmp_path, None
"postgresql://app@db/litellm", PgBouncerSettings(enabled=True, port=port), socket_dir, None
)
assert isinstance(plan, PgBouncerPlan), plan
ini_path: Final = write_pgbouncer_ini(plan, tmp_path, None)
write_userlist(plan.userlist("first"), tmp_path, None)
ini_path: Final = write_pgbouncer_ini(plan, socket_dir, None)
write_userlist(plan.userlist("first"), socket_dir, None)
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, auth_log=auth_log)), str(ini_path)),
argv=(str(_fake_pooler(socket_dir, port, auth_log=auth_log)), str(ini_path)),
port=port,
socket_path=unix_socket_path(tmp_path, port),
socket_path=unix_socket_path(socket_dir, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
install_pgbouncer_token(plan, tmp_path, None, pooler, "second")
install_pgbouncer_token(plan, socket_dir, None, pooler, "second")
assert _wait_until(lambda: auth_log.read_text().count("\n") == 2)
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
@ -837,9 +839,9 @@ class TestStartInContainerPgBouncer:
assert "1.19" in outcome.reason
assert not _listening(port)
def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, tmp_path: Path):
def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, socket_tempdir: Path):
port: Final = _free_port()
binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.19.0")
binary: Final = _fake_pooler(socket_tempdir, port, version_banner="PgBouncer 1.19.0")
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm")
assert isinstance(pooled, str), pooled

5
tests/unit/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -3,6 +3,7 @@ import json
import logging
from collections.abc import Mapping
from datetime import datetime, timedelta, timezone
from pathlib import Path
from types import SimpleNamespace
from typing import Final, cast
from unittest.mock import AsyncMock, MagicMock, create_autospec, patch
@ -3062,10 +3063,10 @@ async def _no_fallback(line: bytes) -> None:
@pytest.mark.asyncio
async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_skips_the_pipeline(tmp_path):
async def test_async_log_success_event_hands_the_sidecar_a_compact_event_and_skips_the_pipeline(socket_path: Path):
handler = _RecordingHandler()
consumer = SpendEventConsumer(handler)
address = UnixAddress(path=str(tmp_path / "spend.sock"))
address = UnixAddress(path=str(socket_path))
server = await consumer.serve(address)
producer = SpendEventProducer(
address=address, on_unavailable="fallback", buffer_size=10, connect_timeout=1.0, fallback=_no_fallback

40
tests/unit/proxy/spend_tracking/test_spend_event_producer.py
View file

@ -177,9 +177,9 @@ def test_settings_read_the_documented_env(monkeypatch):
@pytest.mark.asyncio
async def test_events_reach_the_sidecar_once_and_in_order(tmp_path: Path):
async def test_events_reach_the_sidecar_once_and_in_order(socket_path: Path):
fallback: Final = _Fallback()
async with _Sidecar(tmp_path / "spend.sock") as sidecar:
async with _Sidecar(socket_path) as sidecar:
producer: Final = _producer(sidecar.path, fallback)
outcomes: Final = [await producer.publish(f"event-{i}\n".encode()) for i in range(20)]
await producer.close(drain_timeout=5.0)
@ -193,9 +193,9 @@ async def test_events_reach_the_sidecar_once_and_in_order(tmp_path: Path):
@pytest.mark.asyncio
async def test_unreachable_sidecar_falls_back_in_process_and_backs_off(tmp_path: Path):
async def test_unreachable_sidecar_falls_back_in_process_and_backs_off(socket_path: Path):
fallback: Final = _Fallback()
producer: Final = _producer(tmp_path / "missing.sock", fallback)
producer: Final = _producer(socket_path.with_name("missing.sock"), fallback)
first: Final = await producer.publish(b"event-1\n")
await asyncio.sleep(0.05)
second: Final = await producer.publish(b"event-2\n")
@ -210,11 +210,11 @@ async def test_unreachable_sidecar_falls_back_in_process_and_backs_off(tmp_path:
@pytest.mark.parametrize("loop_factory", [asyncio.new_event_loop, uvloop.new_event_loop], ids=["asyncio", "uvloop"])
def test_sidecar_hang_up_falls_back_instead_of_losing_events(
tmp_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop]
socket_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop]
):
async def scenario() -> tuple[list[bytes], list[bytes], tuple[int, int, int]]:
fallback: Final = _Fallback()
sidecar: Final = _Sidecar(tmp_path / "spend.sock")
sidecar: Final = _Sidecar(socket_path)
async with sidecar:
producer: Final = _producer(sidecar.path, fallback)
await producer.publish(b"event-1\n")
@ -236,7 +236,7 @@ def test_sidecar_hang_up_falls_back_instead_of_losing_events(
@pytest.mark.parametrize("loop_factory", [asyncio.new_event_loop, uvloop.new_event_loop], ids=["asyncio", "uvloop"])
def test_mid_stream_crash_never_bills_an_event_on_both_sides(
tmp_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop]
socket_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop]
):
"""Events large enough to straddle the kernel buffer, a sidecar that reads some and then drops the socket: a
failed write may only fall back when the sidecar cannot have read the whole line."""
@ -244,7 +244,7 @@ def test_mid_stream_crash_never_bills_an_event_on_both_sides(
async def scenario() -> tuple[list[bytes], list[bytes], tuple[int, int, int]]:
fallback: Final = _Fallback()
sidecar: Final = _CrashingSidecar(tmp_path / "spend.sock", lines_before_crash=3)
sidecar: Final = _CrashingSidecar(socket_path, lines_before_crash=3)
async with sidecar:
producer: Final = _producer(sidecar.path, fallback)
for event in events:
@ -266,11 +266,11 @@ def test_mid_stream_crash_never_bills_an_event_on_both_sides(
@pytest.mark.asyncio
async def test_drain_timeout_hands_the_in_flight_event_to_fallback(tmp_path: Path):
async def test_drain_timeout_hands_the_in_flight_event_to_fallback(socket_path: Path):
"""A sidecar that stops reading leaves one event half-written; cancelling the writer must not lose it."""
fallback: Final = _Fallback()
stuck: Final = b"x" * (4 * 1024 * 1024) + b"\n"
async with _Sidecar(tmp_path / "spend.sock", reads=False) as sidecar:
async with _Sidecar(socket_path, reads=False) as sidecar:
producer: Final = _producer(sidecar.path, fallback)
assert await producer.publish(stuck) == "queued"
await asyncio.sleep(0.1)
@ -282,10 +282,10 @@ async def test_drain_timeout_hands_the_in_flight_event_to_fallback(tmp_path: Pat
@pytest.mark.asyncio
async def test_shutdown_lets_the_writer_finish_a_fallback_already_in_progress(tmp_path: Path):
async def test_shutdown_lets_the_writer_finish_a_fallback_already_in_progress(socket_path: Path):
"""Cancelling the writer while it runs the pipeline in-process must neither lose nor repeat that event."""
fallback: Final = _GatedFallback()
producer: Final = _producer(tmp_path / "missing.sock", fallback)
producer: Final = _producer(socket_path.with_name("missing.sock"), fallback)
assert await producer.publish(b"event-1\n") == "queued"
await asyncio.wait_for(fallback.started.wait(), 5.0)
closing: Final = asyncio.ensure_future(producer.close(drain_timeout=0.05))
@ -299,10 +299,10 @@ async def test_shutdown_lets_the_writer_finish_a_fallback_already_in_progress(tm
@pytest.mark.asyncio
async def test_shutdown_does_not_replay_an_event_the_kernel_already_took(tmp_path: Path):
async def test_shutdown_does_not_replay_an_event_the_kernel_already_took(socket_path: Path):
"""Cancelling a drain whose bytes already left the process must not run the event a second time in-process."""
fallback: Final = _Fallback()
async with _Sidecar(tmp_path / "spend.sock") as sidecar:
async with _Sidecar(socket_path) as sidecar:
producer: Final = _producer(sidecar.path, fallback, open_connection=_open_with_stalled_drain)
assert await producer.publish(b"event-1\n") == "queued"
await asyncio.sleep(0.1)
@ -316,9 +316,9 @@ async def test_shutdown_does_not_replay_an_event_the_kernel_already_took(tmp_pat
@pytest.mark.asyncio
async def test_drop_policy_counts_instead_of_running_in_process(tmp_path: Path):
async def test_drop_policy_counts_instead_of_running_in_process(socket_path: Path):
fallback: Final = _Fallback()
producer: Final = _producer(tmp_path / "missing.sock", fallback, on_unavailable="drop")
producer: Final = _producer(socket_path.with_name("missing.sock"), fallback, on_unavailable="drop")
await producer.publish(b"event-1\n")
await producer.close(drain_timeout=5.0)
assert await producer.publish(b"event-2\n") == "dropped"
@ -328,9 +328,9 @@ async def test_drop_policy_counts_instead_of_running_in_process(tmp_path: Path):
@pytest.mark.asyncio
async def test_full_buffer_applies_the_unavailable_policy_immediately(tmp_path: Path):
async def test_full_buffer_applies_the_unavailable_policy_immediately(socket_path: Path):
fallback: Final = _Fallback()
async with _Sidecar(tmp_path / "spend.sock") as sidecar:
async with _Sidecar(socket_path) as sidecar:
producer: Final = _producer(sidecar.path, fallback, buffer_size=2)
outcomes: Final = [await producer.publish(f"event-{i}\n".encode()) for i in range(3)]
await producer.close(drain_timeout=5.0)
@ -342,9 +342,9 @@ async def test_full_buffer_applies_the_unavailable_policy_immediately(tmp_path:
@pytest.mark.asyncio
async def test_close_flushes_buffered_events_then_refuses_new_ones(tmp_path: Path):
async def test_close_flushes_buffered_events_then_refuses_new_ones(socket_path: Path):
fallback: Final = _Fallback()
async with _Sidecar(tmp_path / "spend.sock") as sidecar:
async with _Sidecar(socket_path) as sidecar:
producer: Final = _producer(sidecar.path, fallback)
for i in range(50):
await producer.publish(f"event-{i}\n".encode())

22
tests/unit/proxy/test_collector.py
View file

@ -49,14 +49,14 @@ class _Fallback:
@pytest.mark.asyncio
@pytest.mark.parametrize("transport", ["unix", "tcp"])
async def test_consumer_handles_each_producer_line_once_in_order(tmp_path: Path, transport: str):
async def test_consumer_handles_each_producer_line_once_in_order(socket_path: Path, transport: str):
handler: Final = _Handler(fail_on=b"event-3\n")
consumer: Final = SpendEventConsumer(handler)
server: Final = await consumer.serve(
UnixAddress(path=str(tmp_path / "spend.sock")) if transport == "unix" else TcpAddress("127.0.0.1", 0)
UnixAddress(path=str(socket_path)) if transport == "unix" else TcpAddress("127.0.0.1", 0)
)
address: Final = (
UnixAddress(path=str(tmp_path / "spend.sock"))
UnixAddress(path=str(socket_path))
if transport == "unix"
else TcpAddress("127.0.0.1", server.sockets[0].getsockname()[1])
)
@ -74,10 +74,10 @@ async def test_consumer_handles_each_producer_line_once_in_order(tmp_path: Path,
@pytest.mark.asyncio
async def test_consumer_discards_a_truncated_trailing_event(tmp_path: Path):
async def test_consumer_discards_a_truncated_trailing_event(socket_path: Path):
handler: Final = _Handler()
consumer: Final = SpendEventConsumer(handler)
address: Final = UnixAddress(path=str(tmp_path / "spend.sock"))
address: Final = UnixAddress(path=str(socket_path))
server: Final = await consumer.serve(address)
_, writer = await open_collector_connection(address, timeout=1.0)
writer.write(b"whole\npartial-without-newline")
@ -93,9 +93,9 @@ async def test_consumer_discards_a_truncated_trailing_event(tmp_path: Path):
@pytest.mark.asyncio
async def test_drain_reports_producers_still_connected_after_the_timeout(tmp_path: Path):
async def test_drain_reports_producers_still_connected_after_the_timeout(socket_path: Path):
consumer: Final = SpendEventConsumer(_Handler())
address: Final = UnixAddress(path=str(tmp_path / "spend.sock"))
address: Final = UnixAddress(path=str(socket_path))
server: Final = await consumer.serve(address)
_, writer = await open_collector_connection(address, timeout=1.0)
await asyncio.sleep(0.05)
@ -108,11 +108,11 @@ async def test_drain_reports_producers_still_connected_after_the_timeout(tmp_pat
@pytest.mark.asyncio
async def test_graceful_stop_hands_the_producer_over_to_its_fallback_without_losing_events(tmp_path: Path):
async def test_graceful_stop_hands_the_producer_over_to_its_fallback_without_losing_events(socket_path: Path):
handler: Final = _Handler()
fallback: Final = _Fallback()
consumer: Final = SpendEventConsumer(handler)
address: Final = UnixAddress(path=str(tmp_path / "spend.sock"))
address: Final = UnixAddress(path=str(socket_path))
server: Final = await consumer.serve(address)
producer: Final = SpendEventProducer(
address=address, on_unavailable="fallback", buffer_size=100, connect_timeout=1.0, fallback=fallback
@ -134,7 +134,7 @@ async def test_graceful_stop_hands_the_producer_over_to_its_fallback_without_los
@pytest.mark.parametrize("loop_factory", [asyncio.new_event_loop, uvloop.new_event_loop], ids=["asyncio", "uvloop"])
def test_drain_still_hands_over_live_producers_when_another_connection_already_died(
tmp_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop]
socket_path: Path, loop_factory: Callable[[], asyncio.AbstractEventLoop]
):
"""A transport the loop force-closed under a busy handler must not abort the half-close of the others."""
@ -145,7 +145,7 @@ def test_drain_still_hands_over_live_producers_when_another_connection_already_d
await release.wait()
consumer: Final = SpendEventConsumer(slow_handler)
address: Final = UnixAddress(path=str(tmp_path / "spend.sock"))
address: Final = UnixAddress(path=str(socket_path))
server: Final = await consumer.serve(address)
_, dead = await open_collector_connection(address, timeout=1.0)
dead.write(b"stuck\n")