mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
797353f13a
Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
import asyncio
|
|
from collections.abc import Iterator
|
|
from typing import Final
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
from prisma.errors import PrismaError
|
|
|
|
from litellm._service_logger import ServiceTypes
|
|
from litellm.proxy.db.db_span import db_span
|
|
from litellm.proxy.db.log_db_metrics import record_db_io
|
|
|
|
|
|
@pytest.fixture
|
|
def service_hooks() -> Iterator[tuple[AsyncMock, AsyncMock]]:
|
|
success: Final = AsyncMock()
|
|
failure: Final = AsyncMock()
|
|
service_logging: Final = MagicMock(async_service_success_hook=success, async_service_failure_hook=failure)
|
|
with patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock(service_logging_obj=service_logging)):
|
|
yield success, failure
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_completed_write_emits_one_db_event_named_for_the_call_and_table(
|
|
service_hooks: tuple[AsyncMock, AsyncMock],
|
|
) -> None:
|
|
success, failure = service_hooks
|
|
|
|
async with db_span("commit_spend_updates", "LiteLLM_UserTable"):
|
|
record_db_io()
|
|
await asyncio.sleep(0)
|
|
|
|
event: Final = success.await_args.kwargs
|
|
assert (event["service"], event["call_type"], event["event_metadata"]) == (
|
|
ServiceTypes.DB,
|
|
"commit_spend_updates",
|
|
{"table_name": "LiteLLM_UserTable"},
|
|
)
|
|
assert event["duration"] == pytest.approx((event["end_time"] - event["start_time"]).total_seconds())
|
|
assert failure.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_prisma_error_inside_the_write_emits_a_db_failure_event_and_propagates(
|
|
service_hooks: tuple[AsyncMock, AsyncMock],
|
|
) -> None:
|
|
success, failure = service_hooks
|
|
|
|
with pytest.raises(PrismaError):
|
|
async with db_span("insert_spend_logs", "LiteLLM_SpendLogs"):
|
|
raise PrismaError("connection reset")
|
|
await asyncio.sleep(0)
|
|
|
|
event: Final = failure.await_args.kwargs
|
|
assert (event["service"], event["call_type"], event["event_metadata"], str(event["error"])) == (
|
|
ServiceTypes.DB,
|
|
"insert_spend_logs",
|
|
{"table_name": "LiteLLM_SpendLogs"},
|
|
"connection reset",
|
|
)
|
|
assert success.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_dropped_query_engine_connection_emits_a_db_failure_event(
|
|
service_hooks: tuple[AsyncMock, AsyncMock],
|
|
) -> None:
|
|
success, failure = service_hooks
|
|
|
|
with pytest.raises(httpx.ReadError):
|
|
async with db_span("write_tool_spend", "LiteLLM_DailyToolSpend"):
|
|
raise httpx.ReadError("peer closed connection")
|
|
await asyncio.sleep(0)
|
|
|
|
event: Final = failure.await_args.kwargs
|
|
assert (event["call_type"], event["event_metadata"], str(event["error"])) == (
|
|
"write_tool_spend",
|
|
{"table_name": "LiteLLM_DailyToolSpend"},
|
|
"peer closed connection",
|
|
)
|
|
assert success.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_non_database_error_inside_the_write_emits_no_db_event(
|
|
service_hooks: tuple[AsyncMock, AsyncMock],
|
|
) -> None:
|
|
success, failure = service_hooks
|
|
|
|
with pytest.raises(ValueError, match="bad row"):
|
|
async with db_span("insert_spend_logs", "LiteLLM_SpendLogs"):
|
|
raise ValueError("bad row")
|
|
await asyncio.sleep(0)
|
|
|
|
assert (success.await_count, failure.await_count) == (0, 0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_raising_failure_hook_never_replaces_the_prisma_error(
|
|
service_hooks: tuple[AsyncMock, AsyncMock],
|
|
) -> None:
|
|
success, failure = service_hooks
|
|
failure.side_effect = RuntimeError("exporter down")
|
|
|
|
with pytest.raises(PrismaError):
|
|
async with db_span("commit_spend_updates", "LiteLLM_UserTable"):
|
|
raise PrismaError("connection reset")
|
|
|
|
assert failure.await_count == 1
|
|
assert success.await_count == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_block_whose_prisma_client_never_reached_the_engine_emits_no_db_event(
|
|
service_hooks: tuple[AsyncMock, AsyncMock],
|
|
) -> None:
|
|
success, failure = service_hooks
|
|
|
|
async with db_span("team_user_spend", "LiteLLM_SpendLogs"):
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
|
|
assert (success.await_count, failure.await_count) == (0, 0)
|