mirror of
https://github.com/tiennm99/litellm.git
synced 2026-06-18 09:32:08 +00:00
Yuta Saito
58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
"""Tests for LiteLLM CloudZero database helper."""
|
|
|
|
from datetime import datetime, timezone
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from litellm.integrations.cloudzero.database import LiteLLMDatabase
|
|
|
|
|
|
def _setup_db(monkeypatch: pytest.MonkeyPatch, query_return):
|
|
"""Return a database instance with prisma client mocked out."""
|
|
query_mock = AsyncMock(return_value=query_return)
|
|
mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=query_mock))
|
|
db = LiteLLMDatabase()
|
|
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
|
|
return db, query_mock
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_usage_data_parameterized(monkeypatch: pytest.MonkeyPatch):
|
|
"""Start/end filters and limit should be parameterized via placeholders."""
|
|
start = datetime(2024, 5, 1, tzinfo=timezone.utc)
|
|
end = datetime(2024, 5, 2, tzinfo=timezone.utc)
|
|
db, query_mock = _setup_db(monkeypatch, [])
|
|
|
|
await db.get_usage_data(limit=10, 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, 10]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_usage_data_handles_missing_filters(monkeypatch: pytest.MonkeyPatch):
|
|
"""When no filters provided the params should be None placeholders."""
|
|
db, query_mock = _setup_db(monkeypatch, [])
|
|
|
|
await db.get_usage_data()
|
|
|
|
query_text, *params = query_mock.await_args.args
|
|
assert "LIMIT $3" not in query_text
|
|
assert params == [None, None]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_get_usage_data_rejects_invalid_limit(monkeypatch: pytest.MonkeyPatch):
|
|
"""limit must coerce to int or raise ValueError before hitting the DB."""
|
|
db, query_mock = _setup_db(monkeypatch, [])
|
|
|
|
with pytest.raises(ValueError):
|
|
await db.get_usage_data(limit="invalid")
|
|
|
|
assert query_mock.await_count == 0
|