mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
24584d3d3d
* test(proxy): move proxy_server, _experimental and db tests into tests/unit/proxy Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): keep tuple identity in proxy state restore and fix misc target paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yuneng <yuneng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
85 lines
2.9 KiB
Python
85 lines
2.9 KiB
Python
"""The opt-in REPLICA IDENTITY FULL step, without a database.
|
|
|
|
The behavior against real Postgres is covered by
|
|
tests/proxy_migration_tests/test_replica_identity_full.py; these pin the two
|
|
things that hold with no database at all: the statement handed to the Prisma
|
|
CLI, and the promise that no failure of this optional step escapes into a
|
|
migration run that already succeeded.
|
|
"""
|
|
|
|
import subprocess
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from litellm_proxy_extras.replica_identity import (
|
|
REPLICA_IDENTITY_FULL_ENV_VAR,
|
|
apply_replica_identity_full,
|
|
)
|
|
from litellm_proxy_extras.utils import ProxyExtrasDBManager
|
|
|
|
|
|
def test_hands_the_alter_statement_to_the_prisma_cli():
|
|
captured = {}
|
|
|
|
def capture(cmd, **kwargs):
|
|
captured["cmd"] = cmd
|
|
captured["sql"] = Path(cmd[cmd.index("--file") + 1]).read_text()
|
|
return subprocess.CompletedProcess(cmd, 0)
|
|
|
|
with patch(
|
|
"litellm_proxy_extras.replica_identity.run_prisma", side_effect=capture
|
|
):
|
|
applied = apply_replica_identity_full(
|
|
schema_path="/somewhere/schema.prisma",
|
|
prisma_command="prisma",
|
|
prisma_env={"DATABASE_URL": "postgresql://x/y"},
|
|
)
|
|
|
|
assert applied is True
|
|
assert captured["cmd"][:3] == ["prisma", "db", "execute"]
|
|
assert captured["cmd"][-2:] == ["--schema", "/somewhere/schema.prisma"]
|
|
|
|
sql = captured["sql"]
|
|
assert "ALTER TABLE %s REPLICA IDENTITY FULL" in sql
|
|
assert r"c.relname LIKE 'LiteLLM\_%'" in sql
|
|
assert "c.relreplident <> 'f'" in sql
|
|
assert "lock_timeout" in sql
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"failure",
|
|
[
|
|
subprocess.CalledProcessError(1, "prisma", stderr="must be owner of table"),
|
|
subprocess.TimeoutExpired("prisma", 60),
|
|
OSError(2, "No such file or directory"),
|
|
PermissionError(13, "Read-only file system"),
|
|
],
|
|
ids=["rejected", "timed-out", "cli-missing", "read-only-fs"],
|
|
)
|
|
def test_every_failure_is_reported_instead_of_raised(failure):
|
|
with patch(
|
|
"litellm_proxy_extras.replica_identity.run_prisma", side_effect=failure
|
|
):
|
|
assert (
|
|
apply_replica_identity_full(
|
|
schema_path="/somewhere/schema.prisma",
|
|
prisma_command="prisma",
|
|
prisma_env={},
|
|
)
|
|
is False
|
|
)
|
|
|
|
|
|
def test_an_unusable_migrations_dir_skips_the_step_instead_of_killing_the_run(
|
|
tmp_path, monkeypatch
|
|
):
|
|
"""LITELLM_MIGRATION_DIR makes the step copy the migrations tree before it
|
|
can run, and that copy is filesystem work that can fail on its own."""
|
|
blocker = tmp_path / "blocker"
|
|
blocker.write_text("not a directory")
|
|
monkeypatch.setenv(REPLICA_IDENTITY_FULL_ENV_VAR, "true")
|
|
monkeypatch.setenv("LITELLM_MIGRATION_DIR", str(blocker / "migrations"))
|
|
|
|
assert ProxyExtrasDBManager.apply_replica_identity_full_if_requested() is False
|