mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
"""Tests for FocusLiteLLMDatabase query construction."""
|
|
|
|
import hashlib
|
|
from datetime import datetime, timezone
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from litellm.integrations.focus.database import FocusLiteLLMDatabase
|
|
|
|
|
|
def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return):
|
|
"""Create a database instance with a stubbed prisma client."""
|
|
query_mock = AsyncMock(return_value=query_return)
|
|
mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=query_mock))
|
|
db = FocusLiteLLMDatabase()
|
|
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
|
|
return db, query_mock
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_parameterize_filters_and_limit(monkeypatch: pytest.MonkeyPatch):
|
|
start = datetime(2024, 1, 1, tzinfo=timezone.utc)
|
|
end = datetime(2024, 1, 2, tzinfo=timezone.utc)
|
|
db, query_mock = _setup_db(monkeypatch, [])
|
|
|
|
await db.get_usage_data(limit=25, start_time_utc=start, end_time_utc=end)
|
|
|
|
query_text, *params = query_mock.await_args.args
|
|
assert "dus.updated_at >= $1::timestamptz" in query_text
|
|
assert "dus.updated_at <= $2::timestamptz" in query_text
|
|
assert "LIMIT $3" in query_text
|
|
assert params == [start, end, 25]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_execute_without_filters(monkeypatch: pytest.MonkeyPatch):
|
|
row = {
|
|
"id": 1,
|
|
"user_id": "user",
|
|
"date": datetime(2024, 1, 1, tzinfo=timezone.utc),
|
|
}
|
|
db, query_mock = _setup_db(monkeypatch, [row])
|
|
|
|
result = await db.get_usage_data()
|
|
|
|
query_text, *params = query_mock.await_args.args
|
|
assert "WHERE" not in query_text
|
|
assert "LIMIT $" not in query_text
|
|
assert params == []
|
|
assert result.height == 1
|
|
assert result["id"][0] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_accept_string_timestamps(monkeypatch: pytest.MonkeyPatch):
|
|
db, query_mock = _setup_db(monkeypatch, [])
|
|
|
|
start = "2024-02-01T00:00:00+00:00"
|
|
end = "2024-02-02T00:00:00+00:00"
|
|
await db.get_usage_data(start_time_utc=start, end_time_utc=end)
|
|
|
|
_, *params = query_mock.await_args.args
|
|
assert params == [start, end]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_reject_invalid_limit(monkeypatch: pytest.MonkeyPatch):
|
|
db, query_mock = _setup_db(monkeypatch, [])
|
|
|
|
with pytest.raises(ValueError, match='limit must be an integer'):
|
|
await db.get_usage_data(limit="invalid")
|
|
|
|
assert query_mock.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch):
|
|
db, query_mock = _setup_db(monkeypatch, [])
|
|
|
|
await db.get_usage_data()
|
|
|
|
query_text, *_ = query_mock.await_args.args
|
|
assert (
|
|
"COALESCE(vt.organization_id, tt.organization_id) as organization_id"
|
|
in query_text
|
|
)
|
|
assert "ot.organization_alias as organization_alias" in query_text
|
|
assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(monkeypatch: pytest.MonkeyPatch):
|
|
double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
|
|
joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
|
|
dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
|
|
|
|
async def query_raw(query: str, *params):
|
|
if "sha256(" in query:
|
|
return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
|
|
return [joined_row, dirty_row]
|
|
|
|
mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(side_effect=query_raw)))
|
|
db = FocusLiteLLMDatabase()
|
|
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
|
|
|
|
result = await db.get_usage_data()
|
|
|
|
assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
|
|
assert result["team_id"].to_list() == ["team-0", "team-1"]
|