litellm/tests/unit/integrations/focus/test_focus_database.py

"""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"]


@pytest.mark.asyncio
async def test_table_info_returns_the_daily_user_spend_columns(monkeypatch: pytest.MonkeyPatch):
    columns = [
        {"column_name": "id", "data_type": "text", "is_nullable": "NO"},
        {"column_name": "spend", "data_type": "double precision", "is_nullable": "YES"},
    ]
    db, query_mock = _setup_db(monkeypatch, columns)

    info = await db.get_table_info()

    assert info == {"columns": columns, "table_name": "LiteLLM_DailyUserSpend"}
    query_text, *params = query_mock.await_args.args
    assert "WHERE table_name = 'LiteLLM_DailyUserSpend'" in query_text
    assert params == []