From 770fff7058a77ddb399152f75a7e312a95017423 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Thu, 4 Jun 2026 14:56:02 -0700 Subject: [PATCH] test(proxy): stop running real-DB tests in GitHub Actions unit jobs (#29700) * test(proxy): stop running real-DB tests in GitHub Actions unit jobs GitHub Actions unit jobs were spinning up a Postgres service container, but the only active tests that touched it either used the DB incidentally (a cargo-culted prisma_client.connect()) or were genuine integration tests mislabeled as unit. Mock the incidental ones so the proxy-db job needs no container, and move the tests that genuinely need a database (proxy management behavior, master-key-not-persisted, schema-migration sync) to CircleCI, which is already the real-infrastructure lane. * test(proxy): restore no-unexpected-startup-writes canary in master-key test Greptile noted the hash-match assertion no longer catches other unexpected startup writes (a default key, a rotation artifact). The CircleCI job gives each run a fresh DB, so a clean startup must leave the table empty; add that canary back alongside the precise master-key assertion. --- .circleci/config.yml | 120 +++++++++++ .github/workflows/_test-unit-base.yml | 40 +++- .../workflows/_test-unit-services-base.yml | 190 ------------------ .github/workflows/test-unit-proxy-db.yml | 25 +-- .../test-unit-proxy-mgmt-behavior.yml | 34 ---- .github/workflows/test-unit-security.yml | 28 --- .../test_db_schema_migration.py | 70 +++++++ .../test_master_key_not_in_db.py | 46 +++-- .../test_db_schema_migration.py | 87 -------- tests/proxy_unit_tests/test_jwt.py | 1 - tests/proxy_unit_tests/test_proxy_server.py | 36 ++-- 11 files changed, 272 insertions(+), 405 deletions(-) delete mode 100644 .github/workflows/_test-unit-services-base.yml delete mode 100644 .github/workflows/test-unit-proxy-mgmt-behavior.yml delete mode 100644 .github/workflows/test-unit-security.yml create mode 100644 tests/proxy_migration_tests/test_db_schema_migration.py delete mode 100644 tests/proxy_unit_tests/test_db_schema_migration.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 6ee5634f54..599c58a40e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -452,6 +452,120 @@ jobs: - auth_ui_unit_tests_coverage.xml - auth_ui_unit_tests_coverage + proxy_behavior_tests: + docker: + - *python312_image + - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: litellm_test + working_directory: ~/project + environment: + DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" + steps: + - checkout + - setup_google_dns + - install_uv + - run: + name: Install Dependencies + command: | + uv sync --frozen --all-groups --all-extras --python 3.12 + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" + - run: + name: Seed DB schema via prisma db push + command: | + uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss + - run: + name: Generate Prisma Client + command: uv run --no-sync python -m prisma generate + - run: + name: Run proxy management behavior tests + command: | + mkdir -p test-results + uv run --no-sync python -m pytest tests/proxy_behavior \ + -v --junitxml=test-results/junit.xml --durations=10 + no_output_timeout: 15m + - store_test_results: + path: test-results + + proxy_security_tests: + docker: + - *python312_image + - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: litellm_test + working_directory: ~/project + environment: + DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" + steps: + - checkout + - setup_google_dns + - install_uv + - run: + name: Install Dependencies + command: | + uv sync --frozen --all-groups --all-extras --python 3.12 + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" + - run: + name: Seed DB schema via prisma db push + command: | + uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss + - run: + name: Generate Prisma Client + command: uv run --no-sync python -m prisma generate + - run: + name: Run proxy security tests + command: | + mkdir -p test-results + uv run --no-sync python -m pytest tests/proxy_security_tests \ + -v --junitxml=test-results/junit.xml --durations=10 + no_output_timeout: 15m + - store_test_results: + path: test-results + + schema_migration_check: + docker: + - *python312_image + - image: cimg/postgres:16.0@sha256:b125148bc76e8e8eee5eb3ad6020a3a14110a14e8192f1c645128afebe2e2f84 + environment: + POSTGRES_USER: postgres + POSTGRES_PASSWORD: postgres + POSTGRES_DB: litellm_test + working_directory: ~/project + environment: + # An empty database; the test applies every committed migration itself. + DATABASE_URL: "postgresql://postgres:postgres@localhost:5432/litellm_test" + steps: + - checkout + - setup_google_dns + - install_uv + - run: + name: Install Dependencies + command: | + uv sync --frozen --all-groups --all-extras --python 3.12 + - wait_for_service: + url: tcp://localhost:5432 + timeout: "60" + - run: + name: Generate Prisma Client + command: uv run --no-sync python -m prisma generate + - run: + name: Check schema.prisma is in sync with committed migrations + command: | + mkdir -p test-results + uv run --no-sync python -m pytest tests/proxy_migration_tests \ + -v --junitxml=test-results/junit.xml --durations=10 + no_output_timeout: 15m + - store_test_results: + path: test-results + litellm_router_testing: # Runs all tests with the "router" keyword docker: - *python312_image @@ -2643,6 +2757,12 @@ workflows: filters: *main_branches - auth_ui_unit_tests: filters: *main_branches + - proxy_behavior_tests: + filters: *main_branches + - proxy_security_tests: + filters: *main_branches + - schema_migration_check: + filters: *main_branches - build_docker_database_image: filters: *main_branches - e2e_ui_testing: diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 7e91341ac7..a42b2f8f9d 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -27,6 +27,11 @@ on: required: false type: number default: 10 + dist: + description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)" + required: false + type: string + default: "loadscope" artifact-name: description: "Unique name for the coverage artifact (must be unique per run)" required: true @@ -82,18 +87,31 @@ jobs: MAX_FAILURES: ${{ inputs.max-failures }} WORKERS: ${{ inputs.workers }} RERUNS: ${{ inputs.reruns }} + DIST: ${{ inputs.dist }} run: | - uv run --no-sync pytest ${TEST_PATH:?} \ - --tb=short -vv \ - --maxfail="${MAX_FAILURES}" \ - -n "${WORKERS}" \ - --reruns "${RERUNS}" \ - --reruns-delay 1 \ - --dist=loadscope \ - --durations=20 \ - --cov=./litellm \ - --cov-report=xml:coverage.xml \ - --cov-config=pyproject.toml + if [ "${WORKERS}" = "0" ]; then + uv run --no-sync pytest ${TEST_PATH:?} \ + --tb=short -vv \ + --maxfail="${MAX_FAILURES}" \ + --reruns "${RERUNS}" \ + --reruns-delay 1 \ + --durations=20 \ + --cov=./litellm \ + --cov-report=xml:coverage.xml \ + --cov-config=pyproject.toml + else + uv run --no-sync pytest ${TEST_PATH:?} \ + --tb=short -vv \ + --maxfail="${MAX_FAILURES}" \ + -n "${WORKERS}" \ + --reruns "${RERUNS}" \ + --reruns-delay 1 \ + --dist="${DIST}" \ + --durations=20 \ + --cov=./litellm \ + --cov-report=xml:coverage.xml \ + --cov-config=pyproject.toml + fi - name: Save coverage report if: always() diff --git a/.github/workflows/_test-unit-services-base.yml b/.github/workflows/_test-unit-services-base.yml deleted file mode 100644 index 7f973d8caf..0000000000 --- a/.github/workflows/_test-unit-services-base.yml +++ /dev/null @@ -1,190 +0,0 @@ -name: _Unit Test Services Base (Reusable) - -on: - workflow_call: - inputs: - test-path: - description: "Pytest path(s) to run" - required: true - type: string - workers: - description: "Number of pytest-xdist workers (0 = no parallelism)" - required: false - type: number - default: 2 - reruns: - description: "Number of reruns for flaky tests" - required: false - type: number - default: 2 - timeout-minutes: - description: "Job timeout in minutes" - required: false - type: number - default: 20 - max-failures: - description: "Stop after this many failures" - required: false - type: number - default: 10 - enable-postgres: - description: "Start a local Postgres service container and run Prisma migrations" - required: false - type: boolean - default: false - dist: - description: "pytest-xdist distribution mode (loadscope|load|worksteal|loadfile|no)" - required: false - type: string - default: "loadscope" - artifact-name: - description: "Unique name for the coverage artifact (must be unique per run)" - required: false - type: string - default: "run" - -permissions: - contents: read - -# The postgres service container below is spawned per-job on localhost and -# destroyed with the job. Nothing outside the runner can reach it. The -# user/password/database here are not secrets — they're bootstrap values -# for a throwaway container — so we hardcode them instead of attaching -# every matrix shard to a GHA environment just to read three "secrets" -# (which also produces a "temporarily deployed to …" notification on the -# PR timeline per shard per push). -jobs: - run: - name: Run tests - runs-on: ubuntu-latest - timeout-minutes: ${{ inputs.timeout-minutes }} - - services: - postgres: - image: postgres@sha256:705a5d5b5836f3fcba0d02c4d281e6a7dd9ed2dd4078640f08a1e1e9896e097d # postgres:14 - env: - POSTGRES_USER: litellm - POSTGRES_PASSWORD: litellm - POSTGRES_DB: litellm_test - ports: - - 5432:5432 - options: >- - --health-cmd "pg_isready" - --health-interval 10s - --health-timeout 5s - --health-retries 5 - - steps: - - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Set up Python - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Set up uv - uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7 - with: - version: "0.10.9" - - - name: Cache uv dependencies - uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0 - with: - path: | - ~/.cache/uv - .venv - key: ${{ runner.os }}-uv-services-${{ hashFiles('uv.lock') }} - restore-keys: | - ${{ runner.os }}-uv-services- - - - name: Install dependencies - run: | - uv sync --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router - - - name: Generate Prisma client - env: - PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache - run: | - uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma - - - name: Run Prisma migrations - if: ${{ inputs.enable-postgres }} - env: - DATABASE_URL: "postgresql://litellm:litellm@localhost:5432/litellm_test" - run: | - uv run --no-sync prisma db push --schema litellm/proxy/schema.prisma --accept-data-loss - - - name: Run tests - env: - TEST_PATH: ${{ inputs.test-path }} - MAX_FAILURES: ${{ inputs.max-failures }} - WORKERS: ${{ inputs.workers }} - RERUNS: ${{ inputs.reruns }} - DIST: ${{ inputs.dist }} - DATABASE_URL: ${{ inputs.enable-postgres && 'postgresql://litellm:litellm@localhost:5432/litellm_test' || '' }} - run: | - if [ "${WORKERS}" = "0" ]; then - uv run --no-sync pytest ${TEST_PATH:?} \ - --tb=short -vv \ - --maxfail="${MAX_FAILURES}" \ - --reruns "${RERUNS}" \ - --reruns-delay 1 \ - --durations=20 \ - --cov=./litellm \ - --cov-report=xml:coverage.xml \ - --cov-config=pyproject.toml - else - uv run --no-sync pytest ${TEST_PATH:?} \ - --tb=short -vv \ - --maxfail="${MAX_FAILURES}" \ - -n "${WORKERS}" \ - --reruns "${RERUNS}" \ - --reruns-delay 1 \ - --dist="${DIST}" \ - --durations=20 \ - --cov=./litellm \ - --cov-report=xml:coverage.xml \ - --cov-config=pyproject.toml - fi - - - name: Save coverage report - if: always() - uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 - with: - name: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }} - path: coverage.xml - retention-days: 1 - - upload-coverage: - name: Upload coverage to Codecov - needs: run - if: always() - runs-on: ubuntu-latest - permissions: - contents: read - id-token: write - pull-requests: write - - steps: - - name: Checkout code - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - persist-credentials: false - - - name: Download coverage report - uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e # v4.2.1 - with: - pattern: coverage-${{ inputs.artifact-name }}-${{ github.run_id }}-${{ github.run_attempt }} - path: coverage-reports - merge-multiple: true - - - name: Upload to Codecov - uses: codecov/codecov-action@75cd11691c0faa626561e295848008c8a7dddffe # v5.5.4 - with: - use_oidc: true - directory: coverage-reports - root_dir: ${{ github.workspace }} - flags: ${{ inputs.artifact-name }} - fail_ci_if_error: false diff --git a/.github/workflows/test-unit-proxy-db.yml b/.github/workflows/test-unit-proxy-db.yml index 2d4e85630d..2ac9a3b7c1 100644 --- a/.github/workflows/test-unit-proxy-db.yml +++ b/.github/workflows/test-unit-proxy-db.yml @@ -1,9 +1,10 @@ name: "Unit Tests: Proxy DB Operations" -# Uses DATABASE_URL secret — only runs on trusted branches, not PRs. on: - push: - branches: [main, "litellm_**"] + pull_request: + branches: + - main + - litellm_internal_staging permissions: contents: read @@ -30,9 +31,6 @@ concurrency: # xdist balances its 188 parametrized cases across workers instead of # pinning the whole file to one worker (the default --dist=loadscope # behavior for single-file targets). -# * test_db_schema_migration.py is isolated because one test in it -# (test_aaaasschema_migration_check) takes ~170s — by itself it -# determines the shard's wall-clock floor. jobs: # Fast guard — fails the workflow if a test_*.py file under # tests/proxy_unit_tests/ is not referenced by any matrix entry below. @@ -166,18 +164,6 @@ jobs: dist: loadscope timeout: 15 - # ---- db-and-spend: isolate the 170s schema-migration test ---- - # test_db_schema_migration.py has exactly one test, and that test - # is mostly waiting on `prisma migrate deploy` / `prisma migrate - # diff` subprocesses (~170s). It does no CPU-bound Python work - # inside the test. Running with workers=0 (serial, no xdist) - # skips the 4-worker cold-start cost we'd otherwise pay for a - # single test, saving ~4 minutes of wall-clock. - - test-group: schema-migration - test-path: "tests/proxy_unit_tests/test_db_schema_migration.py" - workers: 0 - dist: loadscope - timeout: 15 - test-group: db-and-spend test-path: >- tests/proxy_unit_tests/test_prisma_client_backoff_retry.py @@ -232,12 +218,11 @@ jobs: workers: 4 dist: loadscope timeout: 15 - uses: ./.github/workflows/_test-unit-services-base.yml + uses: ./.github/workflows/_test-unit-base.yml with: test-path: ${{ matrix.test-path }} workers: ${{ matrix.workers }} reruns: 2 timeout-minutes: ${{ matrix.timeout }} - enable-postgres: true dist: ${{ matrix.dist }} artifact-name: proxy-db-${{ matrix.test-group }} diff --git a/.github/workflows/test-unit-proxy-mgmt-behavior.yml b/.github/workflows/test-unit-proxy-mgmt-behavior.yml deleted file mode 100644 index e73997323a..0000000000 --- a/.github/workflows/test-unit-proxy-mgmt-behavior.yml +++ /dev/null @@ -1,34 +0,0 @@ -name: "Unit Tests: Proxy Management-Endpoint Behavior Pinning" - -on: - pull_request: - branches: - - main - - litellm_internal_staging - - litellm_oss_branch - - "litellm_**" - -permissions: - contents: read - id-token: write - pull-requests: write - -concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} - cancel-in-progress: true - -jobs: - proxy-mgmt-behavior: - uses: ./.github/workflows/_test-unit-services-base.yml - with: - test-path: tests/proxy_behavior - # workers=0 (no xdist): the world seed is a single shared Postgres - # state — two xdist workers both call seed_world() and race on the - # ``behavior-pin-budget`` row, producing UniqueViolation + cascading - # missing-membership FK failures. The whole suite is ~7s sequentially, - # so the cost of disabling parallelism here is negligible. - workers: 0 - reruns: 0 - enable-postgres: true - artifact-name: proxy-mgmt-behavior - timeout-minutes: 15 diff --git a/.github/workflows/test-unit-security.yml b/.github/workflows/test-unit-security.yml deleted file mode 100644 index 4ee8989702..0000000000 --- a/.github/workflows/test-unit-security.yml +++ /dev/null @@ -1,28 +0,0 @@ -name: "Unit Tests: Security" - -# Kept push-only (was previously required by DATABASE_URL secret scoping; -# now the postgres credentials are ephemeral localhost values but the -# push-trigger stays to match the proxy-db workflow cadence). -on: - push: - branches: [main, "litellm_**"] - -permissions: - contents: read - id-token: write - pull-requests: write - -concurrency: - group: ${{ github.workflow }}-${{ github.ref }} - cancel-in-progress: true - -jobs: - security: - uses: ./.github/workflows/_test-unit-services-base.yml - with: - test-path: "tests/proxy_security_tests/" - workers: 1 - reruns: 2 - timeout-minutes: 20 - enable-postgres: true - artifact-name: security diff --git a/tests/proxy_migration_tests/test_db_schema_migration.py b/tests/proxy_migration_tests/test_db_schema_migration.py new file mode 100644 index 0000000000..b0d44cd3e1 --- /dev/null +++ b/tests/proxy_migration_tests/test_db_schema_migration.py @@ -0,0 +1,70 @@ +import os +import shutil +import subprocess +import tempfile +from pathlib import Path + +import pytest + + +@pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) +def test_schema_migration_in_sync(): + """Fail if schema.prisma has changes not captured by the committed migrations. + + Applies every committed migration to an empty database, then diffs the result + against schema.prisma. A non-empty diff means the schema was changed without a + matching migration being generated. + """ + db_url = os.environ["DATABASE_URL"] + source_migrations_dir = Path( + "./litellm-proxy-extras/litellm_proxy_extras/migrations" + ) + source_schema_path = Path("./schema.prisma") + + temp_base = Path(tempfile.mkdtemp(prefix="litellm_schema_migration_")) + schema_path = temp_base / "schema.prisma" + migrations_dir = temp_base / "migrations" + + try: + shutil.copy(source_schema_path, schema_path) + shutil.copytree(source_migrations_dir, migrations_dir) + + if not any(migrations_dir.iterdir()): + pytest.fail( + "No existing migrations found. Run `python litellm/ci_cd/baseline_db_migration.py`." + ) + + subprocess.run( + ["prisma", "migrate", "deploy", "--schema", str(schema_path)], + check=True, + env={**os.environ, "DATABASE_URL": db_url}, + ) + + diff = subprocess.run( + [ + "prisma", + "migrate", + "diff", + "--from-url", + db_url, + "--to-schema-datamodel", + str(schema_path), + "--script", + "--exit-code", + ], + capture_output=True, + text=True, + ) + + if diff.returncode == 2: + pytest.fail( + "Schema changes detected that no migration captures. Run " + "`python litellm/ci_cd/run_migration.py `.\n\n" + + diff.stdout + ) + assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}" + finally: + shutil.rmtree(temp_base, ignore_errors=True) diff --git a/tests/proxy_security_tests/test_master_key_not_in_db.py b/tests/proxy_security_tests/test_master_key_not_in_db.py index 36ac1eb3e2..cb6e08d674 100644 --- a/tests/proxy_security_tests/test_master_key_not_in_db.py +++ b/tests/proxy_security_tests/test_master_key_not_in_db.py @@ -1,39 +1,32 @@ import os import pytest from fastapi.testclient import TestClient -from litellm.proxy.proxy_server import app, ProxyLogging +from litellm.proxy.proxy_server import app, ProxyLogging, hash_token from litellm.caching import DualCache +MASTER_KEY = "sk-1234" + @pytest.fixture(autouse=True) def override_env_settings(monkeypatch): - # Set environment variables only for tests using-monkeypatch (function scope by default). - # Use DATABASE_URL from environment (set by CircleCI to local postgres) if "DATABASE_URL" not in os.environ: pytest.fail( - "DATABASE_URL not set - this test requires a local postgres database to be running" + "DATABASE_URL not set - this test requires a postgres database to be running" ) - monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-1234") + monkeypatch.setenv("LITELLM_MASTER_KEY", MASTER_KEY) monkeypatch.setenv("LITELLM_LOG", "DEBUG") @pytest.fixture(scope="module") def test_client(): - """ - This fixture starts up the test client which triggers FastAPI's startup events. - Prisma will connect to the DB using the provided DATABASE_URL. - """ + """Starting the test client triggers FastAPI startup, where Prisma connects to the DB.""" with TestClient(app) as client: yield client @pytest.mark.asyncio async def test_master_key_not_inserted(test_client): - """ - This test ensures that when the app starts (or when you hit the /health endpoint - to trigger startup logic), no unexpected write occurs in the DB. - """ - # Hit an endpoint (like /health) that triggers any startup tasks. + """The master key must never be persisted to the verification-token table on startup.""" response = test_client.get("/health/liveliness") assert response.status_code == 200 @@ -46,13 +39,22 @@ async def test_master_key_not_inserted(test_client): ), ) - # Connect directly to the test database to inspect the data. await prisma_client.connect() - result = await prisma_client.db.litellm_verificationtoken.find_many() - print(result) + stored_tokens = { + row.token + for row in await prisma_client.db.litellm_verificationtoken.find_many() + } - # The expectation is that no token (or unintended record) is added on startup. - assert len(result) == 0, ( - "SECURITY ALERT SECURITY ALERT SECURITY ALERT: Expected no record in the litellm_verificationtoken table. On startup - the master key should NOT be Inserted into the DB." - "We have found keys in the DB. This is unexpected and should not happen." - ) + for leaked in (hash_token(MASTER_KEY), MASTER_KEY): + assert leaked not in stored_tokens, ( + "SECURITY ALERT: the master key was found in the litellm_verificationtoken " + "table. The master key must never be inserted into the DB." + ) + + # Canary against any other unexpected startup write (default key, rotation + # artifact, ...). The job gives each run a fresh DB, so a clean startup must + # leave the table empty; if startup ever legitimately seeds a token, narrow + # this while keeping the master-key assertion above. + assert ( + not stored_tokens + ), f"startup unexpectedly wrote token(s) to litellm_verificationtoken: {stored_tokens}" diff --git a/tests/proxy_unit_tests/test_db_schema_migration.py b/tests/proxy_unit_tests/test_db_schema_migration.py deleted file mode 100644 index bfd46f4b3d..0000000000 --- a/tests/proxy_unit_tests/test_db_schema_migration.py +++ /dev/null @@ -1,87 +0,0 @@ -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 - - -@pytest.mark.xdist_group("proxy_heavy") -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://test-user:test-password@test-host.example.com/test-db?sslmode=require" - monkeypatch.setenv("DATABASE_URL", test_db_url) - - deploy_dir = Path("./litellm-proxy-extras/litellm_proxy_extras") - source_migrations_dir = deploy_dir / "migrations" - source_schema_path = Path("./schema.prisma") - - # Use worker-specific temp directory to avoid races when running with -n 8. - # Prisma expects migrations in /migrations, so we create that layout. - temp_base = Path(tempfile.mkdtemp(prefix="litellm_schema_migration_")) - temp_migrations_dir = temp_base / "migrations" - schema_path = temp_base / "schema.prisma" - - try: - shutil.copy(source_schema_path, schema_path) - 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 `." - ) - else: - print("No schema changes detected. Migration not needed.") - - finally: - # Clean up: remove temporary directory - if temp_base.exists(): - shutil.rmtree(temp_base) diff --git a/tests/proxy_unit_tests/test_jwt.py b/tests/proxy_unit_tests/test_jwt.py index 92209e1131..beaa120dcb 100644 --- a/tests/proxy_unit_tests/test_jwt.py +++ b/tests/proxy_unit_tests/test_jwt.py @@ -747,7 +747,6 @@ async def test_allowed_routes_admin( from litellm.proxy.proxy_server import user_api_key_auth setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) - await litellm.proxy.proxy_server.prisma_client.connect() monkeypatch.setenv("JWT_PUBLIC_KEY_URL", "https://example.com/public-key") diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 6fdabe64e2..e4fca7ceb0 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -32,7 +32,7 @@ logging.basicConfig( format="%(asctime)s - %(levelname)s - %(message)s", ) -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch from fastapi import FastAPI @@ -1123,6 +1123,14 @@ from litellm.proxy.management_endpoints.team_endpoints import team_member_add from test_key_generate_prisma import prisma_client +@pytest.fixture +def mock_prisma_client(): + client = MagicMock() + client.connect = AsyncMock() + client.disconnect = AsyncMock() + return client + + @pytest.mark.skip(reason="Requires reliable external DB connection (prisma).") @pytest.mark.parametrize( "user_role", @@ -1289,7 +1297,6 @@ async def test_create_team_member_add_team_admin_user_api_key_auth( setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm, "max_internal_user_budget", 10) setattr(litellm, "internal_user_budget_duration", "5m") - await litellm.proxy.proxy_server.prisma_client.connect() user = f"ishaan {uuid.uuid4().hex}" _team_id = "litellm-test-client-id-new" user_key = "sk-12345678" @@ -1364,7 +1371,6 @@ async def test_create_team_member_add_team_admin( setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") setattr(litellm, "max_internal_user_budget", 10) setattr(litellm, "internal_user_budget_duration", "5m") - await litellm.proxy.proxy_server.prisma_client.connect() user = f"ishaan {uuid.uuid4().hex}" _team_id = "litellm-test-client-id-new" user_key = "sk-12345678" @@ -1605,7 +1611,10 @@ async def test_add_callback_via_key(prisma_client): ], ) async def test_add_callback_via_key_litellm_pre_call_utils( - prisma_client, callback_type, expected_success_callbacks, expected_failure_callbacks + mock_prisma_client, + callback_type, + expected_success_callbacks, + expected_failure_callbacks, ): import json @@ -1614,9 +1623,8 @@ async def test_add_callback_via_key_litellm_pre_call_utils( from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") - await litellm.proxy.proxy_server.prisma_client.connect() proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") @@ -1762,7 +1770,10 @@ async def test_disable_fallbacks_by_key(disable_fallbacks_set): ], ) async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket( - prisma_client, callback_type, expected_success_callbacks, expected_failure_callbacks + mock_prisma_client, + callback_type, + expected_success_callbacks, + expected_failure_callbacks, ): import json @@ -1771,9 +1782,8 @@ async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket( from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") - await litellm.proxy.proxy_server.prisma_client.connect() proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config") @@ -1896,7 +1906,10 @@ async def test_add_callback_via_key_litellm_pre_call_utils_gcs_bucket( ], ) async def test_add_callback_via_key_litellm_pre_call_utils_langsmith( - prisma_client, callback_type, expected_success_callbacks, expected_failure_callbacks + mock_prisma_client, + callback_type, + expected_success_callbacks, + expected_failure_callbacks, ): import json @@ -1905,9 +1918,8 @@ async def test_add_callback_via_key_litellm_pre_call_utils_langsmith( from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request - setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "prisma_client", mock_prisma_client) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") - await litellm.proxy.proxy_server.prisma_client.connect() proxy_config = getattr(litellm.proxy.proxy_server, "proxy_config")