DBSpendUpdateWriter

This commit is contained in:
Ishaan Jaff 2025-03-27 16:43:18 -07:00
parent 072be44a54
commit a0fd508de4
2 changed files with 7 additions and 5 deletions

View file

@ -7,9 +7,8 @@ from typing import Any, Optional, Union
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import LiteLLM_UserTable, SpendLogsPayload
from litellm.proxy.proxy_server import hash_token
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend
from litellm.proxy.utils import PrismaClient, ProxyUpdateSpend, hash_token
class DBSpendUpdateWriter:

View file

@ -416,7 +416,8 @@ class TestSpendLogsPayload:
# litellm._turn_on_debug()
with patch.object(
litellm.proxy.proxy_server, "_set_spend_logs_payload"
litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter,
"_set_spend_logs_payload",
) as mock_client, patch.object(litellm.proxy.proxy_server, "prisma_client"):
response = await litellm.acompletion(
model="gpt-4o",
@ -509,7 +510,8 @@ class TestSpendLogsPayload:
client = AsyncHTTPHandler()
with patch.object(
litellm.proxy.proxy_server, "_set_spend_logs_payload"
litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter,
"_set_spend_logs_payload",
) as mock_client, patch.object(
litellm.proxy.proxy_server, "prisma_client"
), patch.object(
@ -604,7 +606,8 @@ class TestSpendLogsPayload:
)
with patch.object(
litellm.proxy.proxy_server, "_set_spend_logs_payload"
litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter,
"_set_spend_logs_payload",
) as mock_client, patch.object(
litellm.proxy.proxy_server, "prisma_client"
), patch.object(