ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

test(proxy): stop running real-DB tests in GitHub Actions unit jobs (#29700)

Browse Source 
* 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.
This commit is contained in:
ryan-crabbe-berri 2026-06-04 14:56:02 -07:00 committed by GitHub
parent 1dbf46665e
commit 770fff7058
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
11 changed files with 272 additions and 405 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

120
.circleci/config.yml
View File

@ -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:

40
.github/workflows/_test-unit-base.yml vendored
View File

@ -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()

190
.github/workflows/_test-unit-services-base.yml vendored
View File

@ -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

25
.github/workflows/test-unit-proxy-db.yml vendored
View File

@ -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 }}

34
.github/workflows/test-unit-proxy-mgmt-behavior.yml vendored
View File

@ -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

28
.github/workflows/test-unit-security.yml vendored
View File

@ -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

70
tests/proxy_migration_tests/test_db_schema_migration.py Normal file
View File

@ -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 <migration_name>`.\n\n"
+ diff.stdout
)
assert diff.returncode == 0, f"prisma migrate diff errored: {diff.stderr}"
finally:
shutil.rmtree(temp_base, ignore_errors=True)

46
tests/proxy_security_tests/test_master_key_not_in_db.py
View File

@ -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}"

87
tests/proxy_unit_tests/test_db_schema_migration.py
View File

@ -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 <schema_dir>/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 <migration_name>`."
)
else:
print("No schema changes detected. Migration not needed.")
finally:
# Clean up: remove temporary directory
if temp_base.exists():
shutil.rmtree(temp_base)

1
tests/proxy_unit_tests/test_jwt.py
View File

@ -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")

36
tests/proxy_unit_tests/test_proxy_server.py
View File

@ -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")