mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
99824533ff
* fix(proxy): kill the whole prisma process group when a boot migration command times out * fix(proxy): keep boot alive when the installed proxy extras has no prisma runner
71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
|
|
import pytest
|
|
|
|
|
|
|
|
def test_check_migration_out_of_sync(mocker):
|
|
"""
|
|
Test that the check_prisma_schema_diff function
|
|
- 🚨 [IMPORTANT] Does NOT Raise an Exception when the Prisma schema is out of sync with the database.
|
|
- logs an error when the Prisma schema is out of sync with the database.
|
|
"""
|
|
# Import the module first so check_migration is in sys.modules,
|
|
# then patch the logger reference in that module directly (not the source
|
|
# module) so the patch works regardless of import order or xdist worker
|
|
# assignment.
|
|
from litellm.proxy.db import check_migration
|
|
|
|
# Mock the helper function to simulate out-of-sync state
|
|
mock_logger = mocker.patch.object(
|
|
check_migration,
|
|
"verbose_logger",
|
|
autospec=True,
|
|
)
|
|
mocker.patch.object(
|
|
check_migration,
|
|
"check_prisma_schema_diff_helper",
|
|
return_value=(True, ["ALTER TABLE users ADD COLUMN new_field TEXT;"]),
|
|
)
|
|
|
|
# Run the function - it should not raise an error
|
|
try:
|
|
check_migration.check_prisma_schema_diff(db_url="mock_url")
|
|
except Exception as e:
|
|
pytest.fail(f"check_prisma_schema_diff raised an unexpected exception: {e}")
|
|
|
|
# Verify the logger was called with the expected message
|
|
check_migration.verbose_logger.exception.assert_called_once()
|
|
actual_message = check_migration.verbose_logger.exception.call_args[0][0]
|
|
assert "prisma schema out of sync with db" in actual_message
|
|
|
|
|
|
@pytest.mark.timeout(30)
|
|
def test_migrate_diff_stops_at_its_budget_and_takes_its_process_tree_with_it(fake_prisma_cli, monkeypatch):
|
|
"""
|
|
`prisma migrate diff` ran unbounded, so a database that never answers hung boot
|
|
before uvicorn ever started, and interrupting the proxy orphaned the schema engine.
|
|
"""
|
|
from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
|
|
|
|
monkeypatch.setenv("FAKE_PRISMA_HANG_FIRST", "1")
|
|
|
|
assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])
|
|
assert fake_prisma_cli.calls == [
|
|
["migrate", "diff", "--from-url", "postgresql://u:p@localhost:9/x",
|
|
"--to-schema-datamodel", "./schema.prisma", "--script"]
|
|
]
|
|
assert fake_prisma_cli.grandchild_is_gone(within_seconds=5)
|
|
|
|
|
|
def test_migrate_diff_without_the_prisma_runner_skips_instead_of_crashing_boot(monkeypatch):
|
|
"""
|
|
Boot calls this helper directly, so an ImportError here takes the proxy down before
|
|
uvicorn starts. An install without the runner must lose the diagnostic, not the proxy.
|
|
"""
|
|
import sys
|
|
|
|
from litellm.proxy.db.check_migration import check_prisma_schema_diff_helper
|
|
|
|
monkeypatch.setitem(sys.modules, "litellm_proxy_extras.prisma_toolchain", None)
|
|
|
|
assert check_prisma_schema_diff_helper("postgresql://u:p@localhost:9/x") == (False, [])
|