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
+1 -2
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:
@@ -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(