Spaces:
Configuration error
Configuration error
File size: 3,297 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import pytest
import os
import subprocess
from pathlib import Path
from pytest_postgresql import factories
import shutil
import tempfile
# Create postgresql fixture
postgresql_my_proc = factories.postgresql_proc(port=None)
postgresql_my = factories.postgresql("postgresql_my_proc")
@pytest.fixture(scope="function")
def schema_setup(postgresql_my):
"""Fixture to provide a test postgres database"""
return postgresql_my
def test_aaaasschema_migration_check(schema_setup, monkeypatch):
"""Test to check if schema requires migration"""
# Set test database URL
test_db_url = f"postgresql://{schema_setup.info.user}:@{schema_setup.info.host}:{schema_setup.info.port}/{schema_setup.info.dbname}"
# test_db_url = "postgresql://neondb_owner:npg_JiZPS0DAhRn4@ep-delicate-wave-a55cvbuc.us-east-2.aws.neon.tech/neondb?sslmode=require"
monkeypatch.setenv("DATABASE_URL", test_db_url)
deploy_dir = Path("./litellm-proxy-extras/litellm_proxy_extras")
source_migrations_dir = deploy_dir / "migrations"
schema_path = Path("./schema.prisma")
# Create temporary migrations directory next to schema.prisma
temp_migrations_dir = schema_path.parent / "migrations"
try:
# Copy migrations to correct location
if temp_migrations_dir.exists():
shutil.rmtree(temp_migrations_dir)
shutil.copytree(source_migrations_dir, temp_migrations_dir)
if not temp_migrations_dir.exists() or not any(temp_migrations_dir.iterdir()):
print("No existing migrations found - first migration needed")
pytest.fail(
"No existing migrations found - first migration needed. Run `litellm/ci_cd/baseline_db.py` to create new migration -E.g. `python litellm/ci_cd/baseline_db_migration.py`."
)
# Apply all existing migrations
subprocess.run(
["prisma", "migrate", "deploy", "--schema", str(schema_path)], check=True
)
# Compare current database state against schema
diff_result = subprocess.run(
[
"prisma",
"migrate",
"diff",
"--from-url",
test_db_url,
"--to-schema-datamodel",
str(schema_path),
"--script", # Show the SQL diff
"--exit-code", # Return exit code 2 if there are differences
],
capture_output=True,
text=True,
)
print("Exit code:", diff_result.returncode)
print("Stdout:", diff_result.stdout)
print("Stderr:", diff_result.stderr)
if diff_result.returncode == 2:
print("Schema changes detected. New migration needed.")
print("Schema differences:")
print(diff_result.stdout)
pytest.fail(
"Schema changes detected - new migration required. Run `litellm/ci_cd/run_migration.py` to create new migration -E.g. `python litellm/ci_cd/run_migration.py <migration_name>`."
)
else:
print("No schema changes detected. Migration not needed.")
finally:
# Clean up: remove temporary migrations directory
if temp_migrations_dir.exists():
shutil.rmtree(temp_migrations_dir)
|