scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129
litellm/tests/test_litellm/proxy/test_prisma_engine_watchdog.py
yuneng-jiang 3357ec8d34
test: run the 30 test files stranded in the second mirror (#37595) 
* test: run the 30 test files stranded in the second mirror

tests/litellm sat beside tests/test_litellm, which is the mirror the repo
convention names, and no job collected it. The allowlist called the directory
unresolved and assumed it was a duplicate. It is not: 30 of its 34 files have no
counterpart in the real mirror, so they are tests nobody has run since they were
written, not copies of tests that run elsewhere.

Moving them in is byte-identical, and it is what makes them run. Every one is
now claimed by a shard's test-path rather than by an allowlist entry, and the
216 tests they hold pass. Directories that needed to become packages did, since
several files are named test_transformation.py and pytest cannot import two of
those from non-package directories in one session.

Never running is why three assertions had drifted away from the code:

  * nvidia.nemotron-super-3-120b max_output_tokens, 32000 -> 32768
  * sambanova/MiniMax-M2.7 max_input_tokens, 204800 -> 196608
  * the Vertex text-to-speech handler moved from data= to json=, so the test
    reads the decoded body off the json kwarg instead of parsing the data one

The first two follow model_prices_and_context_window.json, which the catalog
sync keeps current; the third follows the handler. In all three the test was the
stale side.

The lint workflow ran test_no_hardcoded_secrets.py by path and now points at the
new one.

Four files stay behind. Each shares a filename with a live test whose contents
are disjoint from it, so landing those means merging test bodies, which is a
content review rather than a move. The allowlist entry now names those four and
records how many tests each would bring, in place of calling the whole
directory unresolved.

* fix(ci): keep the secret scan out of the mirror's conftest

The secret-scan job runs pytest under uv run --no-project, so its environment
holds pytest and nothing else. That worked while the file sat in tests/litellm,
which has no conftest, and broke the moment it moved into tests/test_litellm,
whose conftest imports litellm on collection: ModuleNotFoundError: No module
named 'dotenv', before a single test ran.

The file is a repo-wide static scan that imports only base64, os, re and pytest,
so it belongs with the other repo-wide checks in tests/code_coverage_tests,
which has no conftest, rather than in the package mirror. Installing the full
dependency set into a 15-second job to satisfy a conftest it does not use would
be the wrong trade.

Verified with the job's exact command:
  uv run --no-project --with 'pytest==9.0.2' pytest \
    tests/code_coverage_tests/test_no_hardcoded_secrets.py -q
  1 passed in 0.47s
2026-08-20 10:59:43 -07:00

568 lines
21 KiB
Python
Raw Permalink Blame History

"""
Tests for PrismaClient engine watchdog: death detection and automatic reconnect.
Covers:
- Engine PID discovery and liveness check
- Engine process gone (os.kill raises ProcessLookupError) → reconnect triggered
- PermissionError from os.kill → treated as alive (process exists but not ours)
- pidfd handler → schedules attempt_db_reconnect even when lock is held
- waitpid thread → instant cross-platform detection, triggers reconnect
- _run_reconnect_cycle branches: heavy path (engine dead) vs lightweight path (engine alive)
- _engine_confirmed_dead flag ensures heavy reconnect even after _engine_pid reset
- Successful heavy reconnect → watcher re-armed for new process
- Missing DATABASE_URL → graceful RuntimeError in reconnect cycle
- Shutdown → polling loop exits cleanly
"""
import asyncio
import os
import threading
import time
from unittest.mock import ANY, AsyncMock, MagicMock, patch
import pytest
from litellm.proxy.utils import PrismaClient, ProxyLogging
@pytest.fixture(autouse=True)
def mock_prisma_binary():
"""Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests."""
import sys
mock_module = MagicMock()
with patch.dict(sys.modules, {"prisma": mock_module}):
yield
@pytest.fixture
def mock_proxy_logging():
proxy_logging = AsyncMock(spec=ProxyLogging)
proxy_logging.failure_handler = AsyncMock()
return proxy_logging
@pytest.fixture
def engine_client(mock_proxy_logging) -> PrismaClient:
"""
Minimal PrismaClient fixture for engine watchdog tests.
Uses the real constructor pattern from PR #21706 (database_url).
"""
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
client.db = MagicMock()
client.db.recreate_prisma_client = AsyncMock()
client.db.disconnect = AsyncMock(return_value=None)
client.db.connect = AsyncMock(return_value=None)
client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
return client
# ---------------------------------------------------------------------------
# _is_engine_alive
# ---------------------------------------------------------------------------
def test_is_engine_alive_returns_true_when_pid_unknown(engine_client):
"""_is_engine_alive returns True when no engine PID is tracked."""
engine_client._engine_pid = 0
assert engine_client._is_engine_alive() is True
def test_is_engine_alive_returns_false_when_process_gone(engine_client):
"""_is_engine_alive returns False when os.kill raises ProcessLookupError."""
engine_client._engine_pid = 9999
with patch("os.kill", side_effect=ProcessLookupError):
assert engine_client._is_engine_alive() is False
def test_is_engine_alive_returns_true_on_permission_error(engine_client):
"""_is_engine_alive returns True when os.kill raises PermissionError (process exists but not ours)."""
engine_client._engine_pid = 1234
with patch("os.kill", side_effect=PermissionError):
assert engine_client._is_engine_alive() is True
def test_is_engine_alive_returns_true_for_running_process(engine_client):
"""_is_engine_alive returns True when os.kill succeeds (process running)."""
engine_client._engine_pid = 1234
with patch("os.kill"):
assert engine_client._is_engine_alive() is True
# ---------------------------------------------------------------------------
# _poll_engine_proc — calls attempt_db_reconnect on death
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_poll_missing_process_triggers_reconnect(engine_client) -> None:
"""Polling loop triggers attempt_db_reconnect when os.kill raises ProcessLookupError."""
engine_client._engine_pid = 1234
engine_client._watching_engine = True
engine_client.attempt_db_reconnect = AsyncMock(return_value=True)
with patch("os.kill", side_effect=ProcessLookupError):
await engine_client._poll_engine_proc()
engine_client.attempt_db_reconnect.assert_awaited_once_with(
reason="engine_process_death",
force=True,
)
@pytest.mark.asyncio
async def test_poll_permission_error_stops_polling(engine_client) -> None:
"""Polling loop stops cleanly when os.kill raises PermissionError (process not ours)."""
engine_client._engine_pid = 1234
engine_client._watching_engine = True
engine_client.attempt_db_reconnect = AsyncMock(return_value=True)
with patch("os.kill", side_effect=PermissionError):
await engine_client._poll_engine_proc()
# PermissionError means process exists but isn't ours — no reconnect, just stop polling
engine_client.attempt_db_reconnect.assert_not_awaited()
assert engine_client._watching_engine is False
assert engine_client._engine_pid == 0
@pytest.mark.asyncio
async def test_stop_loop_halts_polling(engine_client) -> None:
"""Polling loop exits cleanly when _stop_engine_watcher is called."""
engine_client._engine_pid = 1234
engine_client._watching_engine = True
async def stop_during_sleep(_duration: float) -> None:
engine_client._stop_engine_watcher()
with (
patch("os.kill"),
patch("asyncio.sleep", side_effect=stop_during_sleep),
):
await engine_client._poll_engine_proc()
assert engine_client._watching_engine is False
assert engine_client._engine_pid == 0
# ---------------------------------------------------------------------------
# _on_pidfd_readable — calls attempt_db_reconnect
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pidfd_readable_schedules_reconnect(engine_client) -> None:
"""pidfd handler schedules attempt_db_reconnect via asyncio.create_task."""
engine_client._engine_pid = 1234
engine_client.attempt_db_reconnect = AsyncMock(return_value=True)
created_coros = []
def capture_task(coro):
created_coros.append(coro)
return MagicMock()
with patch("asyncio.create_task", side_effect=capture_task):
engine_client._on_pidfd_readable()
# Run the captured coroutine to completion
assert len(created_coros) == 1
await created_coros[0]
engine_client.attempt_db_reconnect.assert_awaited_once_with(
reason="engine_process_death",
force=True,
)
@pytest.mark.asyncio
async def test_pidfd_schedules_reconnect_task_when_lock_held(engine_client) -> None:
"""pidfd handler schedules reconnect task even when _db_reconnect_lock is held."""
engine_client._engine_pid = 1234
created_coros = []
def capture_task(coro):
created_coros.append(coro)
return MagicMock()
async with engine_client._db_reconnect_lock:
with patch("asyncio.create_task", side_effect=capture_task):
engine_client._on_pidfd_readable()
for coro in created_coros:
coro.close()
assert len(created_coros) == 1
# ---------------------------------------------------------------------------
# _run_reconnect_cycle — engine liveness branching
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_heavy_path_when_engine_dead(
engine_client,
) -> None:
"""_run_reconnect_cycle calls recreate_prisma_client when engine is dead."""
engine_client._engine_pid = 1234
engine_client._start_engine_watcher = AsyncMock()
with (
patch.object(engine_client, "_is_engine_alive", return_value=False),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
patch("os.waitpid", side_effect=ChildProcessError),
):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test", expected_generation=ANY
)
engine_client._start_engine_watcher.assert_awaited_once()
engine_client.db.connect.assert_not_awaited()
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead(
engine_client,
) -> None:
"""_run_reconnect_cycle takes heavy path when _engine_confirmed_dead is set.
This is the critical race-condition fix: SIGCHLD/pidfd handlers set
_engine_confirmed_dead BEFORE _cleanup_engine_watcher resets _engine_pid
to 0, so the heavy path executes even after cleanup.
"""
engine_client._engine_pid = 0 # Already reset by cleanup!
engine_client._engine_confirmed_dead = True # But flag survives cleanup
engine_client._start_engine_watcher = AsyncMock()
with (
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
patch("os.waitpid", side_effect=ChildProcessError),
):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.recreate_prisma_client.assert_awaited_once_with(
"postgresql://test", expected_generation=ANY
)
engine_client._start_engine_watcher.assert_awaited_once()
engine_client.db.connect.assert_not_awaited()
assert engine_client._engine_confirmed_dead is False # Reset after use
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_direct_path_when_engine_alive(
engine_client,
) -> None:
"""Direct reconnect (engine alive) probes the writer first and skips the
recreate when the probe is healthy.
The engine-alive path now runs a SELECT 1 probe before recreating. A
healthy probe means the connection is fine — e.g. an IAM token refresh
already replaced the engine (issue #29176) — so recreating would kill a
working engine. Recreate happens only when the probe fails (covered in
test_prisma_client_reconnect.py::
test_run_reconnect_cycle_direct_path_recreates_when_probe_fails). Either
way the blocking `disconnect()` is never called.
"""
engine_client._engine_pid = 1234
engine_client._start_engine_watcher = AsyncMock()
with (
patch.object(engine_client, "_is_engine_alive", return_value=True),
patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}),
):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.disconnect.assert_not_awaited()
engine_client._start_engine_watcher.assert_awaited_once()
@pytest.mark.asyncio
async def test_run_reconnect_cycle_uses_direct_path_when_pid_unknown(
engine_client,
) -> None:
"""When the engine PID is not tracked, direct reconnect still runs and a
healthy probe likewise skips the recreate."""
engine_client._engine_pid = 0
engine_client._start_engine_watcher = AsyncMock()
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.recreate_prisma_client.assert_not_awaited()
engine_client.db.query_raw.assert_awaited_once_with("SELECT 1")
engine_client.db.disconnect.assert_not_awaited()
engine_client._start_engine_watcher.assert_awaited_once()
@pytest.mark.asyncio
async def test_run_reconnect_cycle_heavy_path_raises_without_database_url(
engine_client,
) -> None:
"""Heavy reconnect raises RuntimeError when DATABASE_URL is not set."""
engine_client._engine_pid = 1234
with (
patch.object(engine_client, "_is_engine_alive", return_value=False),
patch.dict(os.environ, {}, clear=True),
patch("os.waitpid", side_effect=ChildProcessError),
):
with pytest.raises(RuntimeError, match="DATABASE_URL not set"):
await engine_client._run_reconnect_cycle(timeout_seconds=5.0)
engine_client.db.recreate_prisma_client.assert_not_awaited()
# ---------------------------------------------------------------------------
# start/stop lifecycle integration
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_start_watchdog_task_also_starts_engine_watcher(
engine_client,
) -> None:
"""start_db_health_watchdog_task() also starts engine watcher."""
engine_client._start_engine_watcher = AsyncMock()
loop = asyncio.get_running_loop()
dummy_task = loop.create_task(asyncio.sleep(3600))
def fake_create_task(coro):
coro.close()
return dummy_task
with patch("asyncio.create_task", side_effect=fake_create_task):
await engine_client.start_db_health_watchdog_task()
engine_client._start_engine_watcher.assert_awaited_once()
dummy_task.cancel()
try:
await dummy_task
except asyncio.CancelledError:
pass
@pytest.mark.asyncio
async def test_stop_watchdog_task_also_stops_engine_watcher(
engine_client,
) -> None:
"""stop_db_health_watchdog_task() also stops engine watcher."""
engine_client._stop_engine_watcher = MagicMock()
loop = asyncio.get_running_loop()
dummy_task = loop.create_task(asyncio.sleep(3600))
engine_client._db_health_watchdog_task = dummy_task
await engine_client.stop_db_health_watchdog_task()
engine_client._stop_engine_watcher.assert_called_once()
assert engine_client._db_health_watchdog_task is None
# ---------------------------------------------------------------------------
# waitpid thread (Unix only; Windows falls back to os.kill polling)
# ---------------------------------------------------------------------------
def test_try_waitpid_watch_returns_false_on_windows(engine_client):
"""_try_waitpid_watch returns False on Windows (os.waitpid/WNOHANG unavailable)."""
with patch("sys.platform", "win32"):
result = engine_client._try_waitpid_watch(1234)
assert result is False
assert engine_client._engine_wait_thread is None
def test_reap_all_zombies_returns_empty_on_windows(engine_client):
"""_reap_all_zombies returns empty set on Windows (waitpid unavailable)."""
with patch("sys.platform", "win32"):
reaped = PrismaClient._reap_all_zombies()
assert reaped == set()
def test_try_waitpid_watch_returns_false_when_not_child(engine_client):
"""_try_waitpid_watch returns False when PID is not our child process."""
engine_client._engine_pid = 9999
with patch("os.waitpid", side_effect=ChildProcessError):
assert engine_client._try_waitpid_watch(9999) is False
assert engine_client._engine_wait_thread is None
def test_try_waitpid_watch_starts_thread_for_child(engine_client):
"""_try_waitpid_watch starts a daemon thread when PID is our child."""
engine_client._engine_pid = 1234
mock_thread = MagicMock()
mock_loop = MagicMock()
with (
patch("os.waitpid", return_value=(0, 0)),
patch("asyncio.get_running_loop", return_value=mock_loop),
patch("threading.Thread", return_value=mock_thread) as mock_thread_cls,
):
result = engine_client._try_waitpid_watch(1234)
assert result is True
mock_thread.start.assert_called_once()
assert engine_client._engine_wait_thread is mock_thread
@pytest.mark.asyncio
async def test_try_waitpid_watch_handles_already_dead_engine(engine_client) -> None:
"""_try_waitpid_watch detects engine already dead at watch start."""
engine_client._engine_pid = 1234
engine_client.attempt_db_reconnect = AsyncMock(return_value=True)
created_coros = []
def capture_task(coro):
created_coros.append(coro)
return MagicMock()
waitpid_calls = iter([(1234, 0)])
def mock_waitpid(pid, flags):
if pid == -1:
raise ChildProcessError
return next(waitpid_calls)
with (
patch("os.waitpid", side_effect=mock_waitpid),
patch("asyncio.create_task", side_effect=capture_task),
):
result = engine_client._try_waitpid_watch(1234)
assert result is True
assert engine_client._engine_confirmed_dead is True
assert len(created_coros) == 1
created_coros[0].close()
@pytest.mark.asyncio
async def test_on_engine_death_from_thread_triggers_reconnect(engine_client) -> None:
"""waitpid thread callback schedules attempt_db_reconnect."""
engine_client._engine_pid = 1234
engine_client.attempt_db_reconnect = AsyncMock(return_value=True)
created_coros = []
def capture_task(coro):
created_coros.append(coro)
return MagicMock()
with patch("asyncio.create_task", side_effect=capture_task):
engine_client._on_engine_death_from_thread(1234)
assert len(created_coros) == 1
await created_coros[0]
engine_client.attempt_db_reconnect.assert_awaited_once_with(
reason="engine_process_death",
force=True,
)
def test_on_engine_death_from_thread_no_double_trigger(engine_client):
"""waitpid thread callback does not trigger reconnect if already confirmed dead."""
engine_client._engine_pid = 1234
engine_client._engine_confirmed_dead = True
with patch("asyncio.create_task") as mock_create_task:
engine_client._on_engine_death_from_thread(1234)
mock_create_task.assert_not_called()
def test_on_engine_death_from_thread_ignores_stale_pid(engine_client):
"""waitpid thread callback ignores death notification for a stale PID."""
engine_client._engine_pid = 5678
with patch("asyncio.create_task") as mock_create_task:
engine_client._on_engine_death_from_thread(1234)
mock_create_task.assert_not_called()
# ---------------------------------------------------------------------------
# Reconnect escalation: lightweight -> heavy after consecutive failures
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_escalation_after_consecutive_direct_reconnect_failures(engine_client):
"""After N consecutive direct reconnect failures, _engine_confirmed_dead
is set to True so _run_reconnect_cycle takes the heavy reconnect path."""
engine_client._reconnect_escalation_threshold = 3
engine_client._consecutive_reconnect_failures = 0
engine_client._db_reconnect_cooldown_seconds = 0 # disable cooldown for test
engine_client._start_engine_watcher = AsyncMock(return_value=None)
# Make the direct path's writer probe fail so it proceeds to recreate
# (a healthy probe would correctly skip recreate), then make recreate
# fail every time.
engine_client.db.query_raw = AsyncMock(side_effect=Exception("probe failed"))
engine_client.db.recreate_prisma_client = AsyncMock(
side_effect=Exception("recreate failed")
)
# Run 3 failed reconnect attempts
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
for _ in range(3):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is False
assert engine_client._consecutive_reconnect_failures == 3
# Next attempt should escalate to the heavy path (recreate_prisma_client still
# the call, but via the _engine_confirmed_dead branch that also re-arms the watcher).
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test_escalation", timeout_seconds=5.0
)
engine_client.db.recreate_prisma_client.assert_awaited_once()
@pytest.mark.asyncio
async def test_successful_reconnect_resets_failure_counter(engine_client):
"""A successful reconnect resets _consecutive_reconnect_failures to 0."""
engine_client._consecutive_reconnect_failures = 2
engine_client._db_reconnect_cooldown_seconds = 0
engine_client._start_engine_watcher = AsyncMock()
# Make reconnect succeed
engine_client.db.recreate_prisma_client = AsyncMock(return_value=None)
engine_client.db.query_raw = AsyncMock(return_value=[{"result": 1}])
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
result = await engine_client._attempt_reconnect_inside_lock(
force=True, reason="test", timeout_seconds=5.0
)
assert result is True
assert engine_client._consecutive_reconnect_failures == 0
def test_escalation_threshold_env_var(mock_proxy_logging):
"""PRISMA_RECONNECT_ESCALATION_THRESHOLD env var is respected."""
with patch.dict(os.environ, {"PRISMA_RECONNECT_ESCALATION_THRESHOLD": "5"}):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
assert client._reconnect_escalation_threshold == 5
def test_escalation_threshold_min_guard(mock_proxy_logging):
"""Escalation threshold cannot be set below 1."""
with patch.dict(os.environ, {"PRISMA_RECONNECT_ESCALATION_THRESHOLD": "0"}):
client = PrismaClient(
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
)
assert client._reconnect_escalation_threshold == 1
Reference in a new issue View git blame Copy permalink