mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
ba6f2b840d
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Waiting to run
LiteLLM Rust DB / db-queries (push) Waiting to run
LiteLLM Rust / rust-lint (push) Waiting to run
LiteLLM Rust / rust-test (push) Waiting to run
LiteLLM Rust / rust-wheel (push) Waiting to run
Terraform Provider / gofmt, vet, build, test (push) Has been cancelled
Terraform Provider / Provider endpoints vs proxy OpenAPI schema (push) Has been cancelled
93 lines
4 KiB
Python
93 lines
4 KiB
Python
import asyncio
|
|
from contextlib import nullcontext
|
|
from typing import Final, LiteralString
|
|
|
|
import pytest
|
|
|
|
from litellm.proxy.common_utils.fips import FipsModeError, enforce_fips_database
|
|
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
|
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
|
from litellm.repositories.fips_repository import FipsRepository
|
|
from tests.integration._support.database import read_rows, scratch_database, write_rows
|
|
|
|
_CONFIG_TABLE: Final = (
|
|
'CREATE TABLE "LiteLLM_Config" (param_name TEXT PRIMARY KEY, param_value JSONB, last_run_at TIMESTAMP(3), '
|
|
"reload_revision BIGINT NOT NULL DEFAULT 0)"
|
|
)
|
|
_KEY_TABLE: Final = 'CREATE TABLE "LiteLLM_VerificationToken" (token TEXT PRIMARY KEY)'
|
|
_MARKER: Final = "INSERT INTO \"LiteLLM_Config\" (param_name, param_value) VALUES ('fips_mode', 'true')"
|
|
_KEY: Final = "INSERT INTO \"LiteLLM_VerificationToken\" (token) VALUES ('hashed-key')"
|
|
_SETTING: Final = "INSERT INTO \"LiteLLM_Config\" (param_name, param_value) VALUES ('general_settings', '{}')"
|
|
|
|
|
|
def _marker_rows(database_url: str) -> int:
|
|
return len(
|
|
read_rows('SELECT 1 FROM "LiteLLM_Config" WHERE param_name = %s', ("fips_mode",), database_url=database_url)
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"fips_mode, setup, refusal, marked_after",
|
|
[
|
|
("true", (_CONFIG_TABLE, _KEY_TABLE), None, 1),
|
|
("true", (_CONFIG_TABLE, _KEY_TABLE, _MARKER, _KEY), None, 1),
|
|
("true", (_CONFIG_TABLE, _KEY_TABLE, _KEY), "already holds data written without it", 0),
|
|
("true", (_CONFIG_TABLE, _KEY_TABLE, _SETTING), "already holds data written without it", 0),
|
|
("true", (), "could not be checked", None),
|
|
("false", (_CONFIG_TABLE, _KEY_TABLE, _MARKER), "set up with LITELLM_FIPS_MODE=true", 1),
|
|
("false", (_CONFIG_TABLE, _KEY_TABLE, _KEY), None, 0),
|
|
("false", (), None, None),
|
|
],
|
|
ids=[
|
|
"fips-on-new-database",
|
|
"fips-on-fips-database",
|
|
"fips-on-database-with-a-key",
|
|
"fips-on-database-with-a-setting",
|
|
"fips-on-database-without-tables",
|
|
"fips-off-fips-database",
|
|
"fips-off-database-with-a-key",
|
|
"fips-off-database-without-tables",
|
|
],
|
|
)
|
|
async def test_fips_database_check(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
fips_mode: str,
|
|
setup: tuple[LiteralString, ...],
|
|
refusal: str | None,
|
|
marked_after: int | None,
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_FIPS_MODE", fips_mode)
|
|
with scratch_database() as database_url:
|
|
for statement in setup:
|
|
write_rows(statement, (), database_url=database_url)
|
|
monkeypatch.setenv("DATABASE_URL", database_url)
|
|
client: Final = PrismaClient(database_url, ProxyLogging(UserApiKeyCache()))
|
|
await client.connect()
|
|
try:
|
|
with pytest.raises(FipsModeError, match=refusal) if refusal else nullcontext():
|
|
await enforce_fips_database(FipsRepository(client), database_is_configured=True)
|
|
finally:
|
|
await client.aclose()
|
|
if marked_after is not None:
|
|
assert _marker_rows(database_url) == marked_after
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pods_starting_together_in_fips_mode_on_a_new_database_both_start(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("LITELLM_FIPS_MODE", "true")
|
|
with scratch_database() as database_url:
|
|
for statement in (_CONFIG_TABLE, _KEY_TABLE):
|
|
write_rows(statement, (), database_url=database_url)
|
|
monkeypatch.setenv("DATABASE_URL", database_url)
|
|
pods: Final = tuple(PrismaClient(database_url, ProxyLogging(UserApiKeyCache())) for _ in range(4))
|
|
await asyncio.gather(*(pod.connect() for pod in pods))
|
|
try:
|
|
await asyncio.gather(
|
|
*(enforce_fips_database(FipsRepository(pod), database_is_configured=True) for pod in pods)
|
|
)
|
|
finally:
|
|
await asyncio.gather(*(pod.aclose() for pod in pods))
|
|
assert _marker_rows(database_url) == 1
|