mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-20 18:24:07 +00:00
Merge pull request #22044 from ryan-crabbe/litellm_redis_pipeline_spend_updates
Litellm redis pipeline spend updates
This commit is contained in:
ryan-crabbe
@@ -22,7 +22,11 @@ from litellm._logging import print_verbose, verbose_logger
|
||||
from litellm.constants import DEFAULT_REDIS_MAJOR_VERSION
|
||||
from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
|
||||
from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
|
||||
from litellm.types.caching import RedisPipelineIncrementOperation
|
||||
from litellm.types.caching import (
|
||||
RedisPipelineIncrementOperation,
|
||||
RedisPipelineLpopOperation,
|
||||
RedisPipelineRpushOperation,
|
||||
)
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
from .base_cache import BaseCache
|
||||
@@ -1320,6 +1324,75 @@ class RedisCache(BaseCache):
|
||||
)
|
||||
raise e
|
||||
|
||||
async def _pipeline_rpush_helper(
|
||||
self,
|
||||
pipe: pipeline,
|
||||
rpush_list: List[RedisPipelineRpushOperation],
|
||||
) -> List[int]:
|
||||
"""Helper function for pipeline rpush operations"""
|
||||
for rpush_op in rpush_list:
|
||||
pipe.rpush(rpush_op["key"], *rpush_op["values"])
|
||||
results = await pipe.execute()
|
||||
# Preserve positional correspondence — raise on per-command errors
|
||||
for r in results:
|
||||
if isinstance(r, Exception):
|
||||
raise r
|
||||
return results
|
||||
|
||||
async def async_rpush_pipeline(
|
||||
self,
|
||||
rpush_list: List[RedisPipelineRpushOperation],
|
||||
) -> List[int]:
|
||||
"""
|
||||
Use Redis Pipelines for bulk RPUSH operations
|
||||
|
||||
Args:
|
||||
rpush_list: List of RedisPipelineRpushOperation dicts containing:
|
||||
- key: str
|
||||
- values: List[Any]
|
||||
|
||||
Returns:
|
||||
List[int]: List lengths after each push
|
||||
"""
|
||||
if len(rpush_list) == 0:
|
||||
return []
|
||||
|
||||
_redis_client: Any = self.init_async_client()
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
async with _redis_client.pipeline(transaction=False) as pipe:
|
||||
results = await self._pipeline_rpush_helper(pipe, rpush_list)
|
||||
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
error=e,
|
||||
call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
)
|
||||
raise e
|
||||
|
||||
async def handle_lpop_count_for_older_redis_versions(
|
||||
self, pipe: pipeline, key: str, count: int
|
||||
) -> List[bytes]:
|
||||
@@ -1400,3 +1473,120 @@ class RedisCache(BaseCache):
|
||||
f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}"
|
||||
)
|
||||
raise e
|
||||
|
||||
async def _pipeline_lpop_helper(
|
||||
self,
|
||||
pipe: pipeline,
|
||||
lpop_list: List[RedisPipelineLpopOperation],
|
||||
) -> List[Optional[List[str]]]:
|
||||
"""Helper function for pipeline lpop operations.
|
||||
|
||||
For Redis >= 7, queues one LPOP(key, count) per operation.
|
||||
For Redis < 7, queues `count` individual LPOP(key) commands per operation.
|
||||
"""
|
||||
major_version = self._parse_redis_major_version()
|
||||
|
||||
if major_version >= 7:
|
||||
for lpop_op in lpop_list:
|
||||
pipe.lpop(lpop_op["key"], lpop_op["count"])
|
||||
raw_results = await pipe.execute()
|
||||
else:
|
||||
# For Redis < 7, LPOP doesn't support count param.
|
||||
# Issue `count` individual LPOP commands per key, all in one pipeline.
|
||||
counts: List[int] = []
|
||||
for lpop_op in lpop_list:
|
||||
count = lpop_op["count"] or 1
|
||||
counts.append(count)
|
||||
for _ in range(count):
|
||||
pipe.lpop(lpop_op["key"])
|
||||
flat_results = await pipe.execute()
|
||||
|
||||
# Re-group the flat results back into per-key lists
|
||||
raw_results = []
|
||||
offset = 0
|
||||
for count in counts:
|
||||
key_results = [
|
||||
r for r in flat_results[offset : offset + count] if r is not None
|
||||
]
|
||||
raw_results.append(key_results if key_results else None)
|
||||
offset += count
|
||||
|
||||
# Raise on per-command errors (matches _pipeline_rpush_helper behavior)
|
||||
for r in raw_results:
|
||||
if isinstance(r, Exception):
|
||||
raise r
|
||||
|
||||
# Decode bytes -> str for each result set
|
||||
decoded_results: List[Optional[List[str]]] = []
|
||||
for r in raw_results:
|
||||
if r is None:
|
||||
decoded_results.append(None)
|
||||
elif isinstance(r, list):
|
||||
try:
|
||||
decoded_results.append(
|
||||
[
|
||||
item.decode("utf-8") if isinstance(item, bytes) else item
|
||||
for item in r
|
||||
if item is not None
|
||||
]
|
||||
or None
|
||||
)
|
||||
except Exception:
|
||||
decoded_results.append(r) # type: ignore
|
||||
else:
|
||||
decoded_results.append(None)
|
||||
return decoded_results
|
||||
|
||||
async def async_lpop_pipeline(
|
||||
self,
|
||||
lpop_list: List[RedisPipelineLpopOperation],
|
||||
) -> List[Optional[List[str]]]:
|
||||
"""
|
||||
Use Redis Pipelines for bulk LPOP operations
|
||||
|
||||
Args:
|
||||
lpop_list: List of RedisPipelineLpopOperation dicts containing:
|
||||
- key: str
|
||||
- count: Optional[int]
|
||||
|
||||
Returns:
|
||||
List[Optional[List[str]]]: Decoded results per key, None if key was empty
|
||||
"""
|
||||
if len(lpop_list) == 0:
|
||||
return []
|
||||
|
||||
_redis_client: Any = self.init_async_client()
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
async with _redis_client.pipeline(transaction=False) as pipe:
|
||||
results = await self._pipeline_lpop_helper(pipe, lpop_list)
|
||||
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_success_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
## LOGGING ##
|
||||
end_time = time.time()
|
||||
_duration = end_time - start_time
|
||||
asyncio.create_task(
|
||||
self.service_logger_obj.async_service_failure_hook(
|
||||
service=ServiceTypes.REDIS,
|
||||
duration=_duration,
|
||||
error=e,
|
||||
call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}",
|
||||
)
|
||||
)
|
||||
verbose_logger.error(
|
||||
"LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s",
|
||||
str(e),
|
||||
)
|
||||
raise e
|
||||
|
||||
@@ -666,9 +666,16 @@ class DBSpendUpdateWriter:
|
||||
verbose_proxy_logger.debug("acquired lock for spend updates")
|
||||
|
||||
try:
|
||||
db_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer()
|
||||
)
|
||||
(
|
||||
db_spend_update_transactions,
|
||||
daily_spend_update_transactions,
|
||||
daily_team_spend_update_transactions,
|
||||
daily_org_spend_update_transactions,
|
||||
daily_end_user_spend_update_transactions,
|
||||
daily_agent_spend_update_transactions,
|
||||
daily_tag_spend_update_transactions,
|
||||
) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
if db_spend_update_transactions is not None:
|
||||
verbose_proxy_logger.info(
|
||||
"Spend tracking - committing spend updates from Redis to DB: "
|
||||
@@ -688,9 +695,6 @@ class DBSpendUpdateWriter:
|
||||
db_spend_update_transactions=db_spend_update_transactions,
|
||||
)
|
||||
|
||||
daily_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_user_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
@@ -698,9 +702,6 @@ class DBSpendUpdateWriter:
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_spend_update_transactions,
|
||||
)
|
||||
daily_team_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_team_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_team_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
@@ -709,9 +710,6 @@ class DBSpendUpdateWriter:
|
||||
daily_spend_transactions=daily_team_spend_update_transactions,
|
||||
)
|
||||
|
||||
daily_org_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_org_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_org_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
@@ -720,9 +718,6 @@ class DBSpendUpdateWriter:
|
||||
daily_spend_transactions=daily_org_spend_update_transactions,
|
||||
)
|
||||
|
||||
daily_tag_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_tag_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_tag_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
@@ -730,9 +725,6 @@ class DBSpendUpdateWriter:
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_tag_spend_update_transactions,
|
||||
)
|
||||
daily_end_user_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_end_user_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_end_user_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
@@ -740,9 +732,6 @@ class DBSpendUpdateWriter:
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
daily_spend_transactions=daily_end_user_spend_update_transactions,
|
||||
)
|
||||
daily_agent_spend_update_transactions = (
|
||||
await self.redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer()
|
||||
)
|
||||
if daily_agent_spend_update_transactions is not None:
|
||||
await DBSpendUpdateWriter.update_daily_agent_spend(
|
||||
n_retry_times=n_retry_times,
|
||||
|
||||
@@ -6,7 +6,7 @@ This is to prevent deadlocks and improve reliability
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching import RedisCache
|
||||
@@ -36,6 +36,7 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
|
||||
)
|
||||
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation
|
||||
from litellm.types.services import ServiceTypes
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -209,47 +210,44 @@ class RedisUpdateBuffer:
|
||||
"ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=db_spend_update_transactions,
|
||||
redis_key=REDIS_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_SPEND_UPDATE_QUEUE,
|
||||
# Build a list of rpush operations, skipping empty/None transaction sets
|
||||
_queue_configs: List[Tuple[Any, str, ServiceTypes]] = [
|
||||
(db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE),
|
||||
(daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE),
|
||||
(daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE),
|
||||
(daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE),
|
||||
(daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE),
|
||||
(daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE),
|
||||
(daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE),
|
||||
]
|
||||
|
||||
rpush_list: List[RedisPipelineRpushOperation] = []
|
||||
service_types: List[ServiceTypes] = []
|
||||
for transactions, redis_key, service_type in _queue_configs:
|
||||
if transactions is None or len(transactions) == 0:
|
||||
continue
|
||||
rpush_list.append(
|
||||
RedisPipelineRpushOperation(
|
||||
key=redis_key,
|
||||
values=[safe_dumps(transactions)],
|
||||
)
|
||||
)
|
||||
service_types.append(service_type)
|
||||
|
||||
if len(rpush_list) == 0:
|
||||
return
|
||||
|
||||
result_lengths = await self.redis_cache.async_rpush_pipeline(
|
||||
rpush_list=rpush_list,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_team_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_org_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_end_user_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_agent_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
|
||||
await self._store_transactions_in_redis(
|
||||
transactions=daily_tag_spend_update_transactions,
|
||||
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
|
||||
service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE,
|
||||
)
|
||||
# Emit gauge events for each queue
|
||||
for i, queue_size in enumerate(result_lengths):
|
||||
if i < len(service_types):
|
||||
await self._emit_new_item_added_to_redis_buffer_event(
|
||||
queue_size=queue_size,
|
||||
service=service_types[i],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _number_of_transactions_to_store_in_redis(
|
||||
@@ -338,6 +336,77 @@ class RedisUpdateBuffer:
|
||||
|
||||
return combined_transaction
|
||||
|
||||
async def get_all_transactions_from_redis_buffer_pipeline(
|
||||
self,
|
||||
) -> Tuple[
|
||||
Optional[DBSpendUpdateTransactions],
|
||||
Optional[Dict[str, DailyUserSpendTransaction]],
|
||||
Optional[Dict[str, DailyTeamSpendTransaction]],
|
||||
Optional[Dict[str, DailyOrganizationSpendTransaction]],
|
||||
Optional[Dict[str, DailyEndUserSpendTransaction]],
|
||||
Optional[Dict[str, DailyAgentSpendTransaction]],
|
||||
Optional[Dict[str, DailyTagSpendTransaction]],
|
||||
]:
|
||||
"""
|
||||
Drains all 7 Redis buffer queues in a single pipeline round-trip.
|
||||
|
||||
Returns a 7-tuple of parsed results in this order:
|
||||
0: DBSpendUpdateTransactions
|
||||
1: daily user spend
|
||||
2: daily team spend
|
||||
3: daily org spend
|
||||
4: daily end-user spend
|
||||
5: daily agent spend
|
||||
6: daily tag spend
|
||||
"""
|
||||
if self.redis_cache is None:
|
||||
return None, None, None, None, None, None, None
|
||||
|
||||
lpop_list: List[RedisPipelineLpopOperation] = [
|
||||
RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
RedisPipelineLpopOperation(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
RedisPipelineLpopOperation(key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
RedisPipelineLpopOperation(key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
RedisPipelineLpopOperation(key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
RedisPipelineLpopOperation(key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
RedisPipelineLpopOperation(key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT),
|
||||
]
|
||||
|
||||
raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
# Pad with None if pipeline returned fewer results than expected
|
||||
while len(raw_results) < 7:
|
||||
raw_results.append(None)
|
||||
|
||||
# Slot 0: DBSpendUpdateTransactions
|
||||
db_spend: Optional[DBSpendUpdateTransactions] = None
|
||||
if raw_results[0] is not None:
|
||||
parsed = self._parse_list_of_transactions(raw_results[0])
|
||||
if len(parsed) > 0:
|
||||
db_spend = self._combine_list_of_transactions(parsed)
|
||||
|
||||
# Slots 1-6: daily spend categories
|
||||
daily_results: List[Optional[Dict[str, Any]]] = []
|
||||
for slot in range(1, 7):
|
||||
if raw_results[slot] is None:
|
||||
daily_results.append(None)
|
||||
else:
|
||||
list_of_daily = [json.loads(t) for t in raw_results[slot]] # type: ignore
|
||||
aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
|
||||
list_of_daily
|
||||
)
|
||||
daily_results.append(aggregated)
|
||||
|
||||
return (
|
||||
db_spend,
|
||||
cast(Optional[Dict[str, DailyUserSpendTransaction]], daily_results[0]),
|
||||
cast(Optional[Dict[str, DailyTeamSpendTransaction]], daily_results[1]),
|
||||
cast(Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2]),
|
||||
cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]),
|
||||
cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]),
|
||||
cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]),
|
||||
)
|
||||
|
||||
async def get_all_daily_spend_update_transactions_from_redis_buffer(
|
||||
self,
|
||||
) -> Optional[Dict[str, DailyUserSpendTransaction]]:
|
||||
|
||||
@@ -52,6 +52,24 @@ class RedisPipelineSetOperation(TypedDict):
|
||||
ttl: Optional[int]
|
||||
|
||||
|
||||
class RedisPipelineRpushOperation(TypedDict):
|
||||
"""
|
||||
TypedDict for 1 Redis Pipeline RPUSH Operation
|
||||
"""
|
||||
|
||||
key: str
|
||||
values: List[Any]
|
||||
|
||||
|
||||
class RedisPipelineLpopOperation(TypedDict):
|
||||
"""
|
||||
TypedDict for 1 Redis Pipeline LPOP Operation
|
||||
"""
|
||||
|
||||
key: str
|
||||
count: Optional[int]
|
||||
|
||||
|
||||
DynamicCacheControl = TypedDict(
|
||||
"DynamicCacheControl",
|
||||
{
|
||||
|
||||
@@ -122,6 +122,249 @@ async def test_handle_lpop_count_for_older_redis_versions(monkeypatch):
|
||||
assert mock_pipeline.execute.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_executes_all_operations(monkeypatch, redis_no_ping):
|
||||
"""Verify that multiple rpush ops are batched into a single pipeline execute"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.rpush = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(return_value=[3, 5, 1])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
rpush_list = [
|
||||
RedisPipelineRpushOperation(key="key1", values=["a", "b"]),
|
||||
RedisPipelineRpushOperation(key="key2", values=["c"]),
|
||||
RedisPipelineRpushOperation(key="key3", values=["d", "e", "f"]),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
result = await redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
|
||||
|
||||
assert result == [3, 5, 1]
|
||||
assert mock_pipeline.rpush.call_count == 3
|
||||
mock_pipeline.rpush.assert_any_call("key1", "a", "b")
|
||||
mock_pipeline.rpush.assert_any_call("key2", "c")
|
||||
mock_pipeline.rpush.assert_any_call("key3", "d", "e", "f")
|
||||
mock_pipeline.execute.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_empty_list_returns_empty(monkeypatch, redis_no_ping):
|
||||
"""Empty rpush_list should return empty list without touching Redis"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
result = await redis_cache.async_rpush_pipeline(rpush_list=[])
|
||||
|
||||
assert result == []
|
||||
mock_redis_instance.pipeline.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_raises_on_redis_error(monkeypatch, redis_no_ping):
|
||||
"""Pipeline errors should propagate"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.rpush = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down"))
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
rpush_list = [RedisPipelineRpushOperation(key="key1", values=["a"])]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
with pytest.raises(ConnectionError, match="Redis down"):
|
||||
await redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_single_round_trip(monkeypatch, redis_no_ping):
|
||||
"""Verify that multiple lpop ops are batched into a single pipeline execute"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
redis_cache.redis_version = "7.0.0"
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.lpop = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(return_value=[
|
||||
[b"val1", b"val2"], # key1 results
|
||||
None, # key2 empty
|
||||
[b"val3"], # key3 results
|
||||
])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineLpopOperation
|
||||
|
||||
lpop_list = [
|
||||
RedisPipelineLpopOperation(key="key1", count=10),
|
||||
RedisPipelineLpopOperation(key="key2", count=10),
|
||||
RedisPipelineLpopOperation(key="key3", count=5),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
assert len(results) == 3
|
||||
assert results[0] == ["val1", "val2"]
|
||||
assert results[1] is None
|
||||
assert results[2] == ["val3"]
|
||||
mock_pipeline.execute.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_redis_lt7_regroups_flat_results(monkeypatch, redis_no_ping):
|
||||
"""Verify Redis < 7 fallback issues individual LPOPs and regroups correctly"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
redis_cache.redis_version = "6.2.0"
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.lpop = MagicMock()
|
||||
|
||||
# With count=3 for key1 and count=2 for key2, we get 5 individual LPOP commands
|
||||
# Simulate: key1 has 2 values then None, key2 has 1 value then None
|
||||
mock_pipeline.execute = AsyncMock(return_value=[
|
||||
b"val1", b"val2", None, # 3 LPOPs for key1
|
||||
b"val3", None, # 2 LPOPs for key2
|
||||
])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineLpopOperation
|
||||
|
||||
lpop_list = [
|
||||
RedisPipelineLpopOperation(key="key1", count=3),
|
||||
RedisPipelineLpopOperation(key="key2", count=2),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0] == ["val1", "val2"] # 2 values, None filtered out
|
||||
assert results[1] == ["val3"] # 1 value, None filtered out
|
||||
# All 5 individual LPOPs should be queued, but only 1 execute() call
|
||||
assert mock_pipeline.lpop.call_count == 5
|
||||
mock_pipeline.execute.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_rpush_pipeline_raises_on_per_command_error(monkeypatch, redis_no_ping):
|
||||
"""Verify that per-command errors in pipeline results are raised, not silently dropped"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.rpush = MagicMock()
|
||||
# Simulate: first RPUSH succeeds, second returns a per-command error
|
||||
mock_pipeline.execute = AsyncMock(return_value=[3, Exception("WRONGTYPE")])
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
rpush_list = [
|
||||
RedisPipelineRpushOperation(key="key1", values=["a"]),
|
||||
RedisPipelineRpushOperation(key="key2", values=["b"]),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
with pytest.raises(Exception, match="WRONGTYPE"):
|
||||
await redis_cache.async_rpush_pipeline(rpush_list=rpush_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_raises_on_per_command_error(monkeypatch, redis_no_ping):
|
||||
"""Verify that per-command errors in LPOP pipeline results are raised, not silently dropped"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
redis_cache.redis_version = "7.0.0"
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.lpop = MagicMock()
|
||||
# Simulate: first LPOP succeeds, second returns a per-command error
|
||||
mock_pipeline.execute = AsyncMock(
|
||||
return_value=[[b"val1"], Exception("WRONGTYPE")]
|
||||
)
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineLpopOperation
|
||||
|
||||
lpop_list = [
|
||||
RedisPipelineLpopOperation(key="key1", count=10),
|
||||
RedisPipelineLpopOperation(key="key2", count=10),
|
||||
]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
with pytest.raises(Exception, match="WRONGTYPE"):
|
||||
await redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping):
|
||||
"""Empty lpop_list should return empty list without touching Redis"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
result = await redis_cache.async_lpop_pipeline(lpop_list=[])
|
||||
|
||||
assert result == []
|
||||
mock_redis_instance.pipeline.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_lpop_pipeline_propagates_redis_exception(monkeypatch, redis_no_ping):
|
||||
"""Pipeline errors should propagate"""
|
||||
monkeypatch.setenv("REDIS_HOST", "https://my-test-host")
|
||||
redis_cache = RedisCache()
|
||||
redis_cache.redis_version = "7.0.0"
|
||||
|
||||
mock_redis_instance = AsyncMock()
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline)
|
||||
mock_pipeline.__aexit__ = AsyncMock(return_value=None)
|
||||
mock_pipeline.lpop = MagicMock()
|
||||
mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down"))
|
||||
mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline)
|
||||
|
||||
from litellm.types.caching import RedisPipelineLpopOperation
|
||||
|
||||
lpop_list = [RedisPipelineLpopOperation(key="key1", count=10)]
|
||||
|
||||
with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance):
|
||||
with pytest.raises(ConnectionError, match="Redis down"):
|
||||
await redis_cache.async_lpop_pipeline(lpop_list=lpop_list)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"redis_version",
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
|
||||
from litellm.types.caching import RedisPipelineRpushOperation
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redis_cache():
|
||||
"""Create a mock RedisCache instance"""
|
||||
mock = AsyncMock()
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def redis_update_buffer(mock_redis_cache):
|
||||
"""Create a RedisUpdateBuffer with a mock RedisCache"""
|
||||
return RedisUpdateBuffer(redis_cache=mock_redis_cache)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache):
|
||||
"""
|
||||
Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once
|
||||
with the correct operations and skips empty queues.
|
||||
"""
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2])
|
||||
|
||||
# Create mock queues - only 3 of 7 have data
|
||||
spend_update_queue = AsyncMock()
|
||||
spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
|
||||
return_value={"key_list_transactions": {"key1": 1.0}}
|
||||
)
|
||||
|
||||
daily_spend_queue = AsyncMock()
|
||||
daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={"user_key1": {"spend": 1.0}}
|
||||
)
|
||||
|
||||
daily_team_queue = AsyncMock()
|
||||
daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={"team_key1": {"spend": 2.0}}
|
||||
)
|
||||
|
||||
# Empty queues
|
||||
daily_org_queue = AsyncMock()
|
||||
daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
daily_end_user_queue = AsyncMock()
|
||||
daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
|
||||
daily_agent_queue = AsyncMock()
|
||||
daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
daily_tag_queue = AsyncMock()
|
||||
daily_tag_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=spend_update_queue,
|
||||
daily_spend_update_queue=daily_spend_queue,
|
||||
daily_team_spend_update_queue=daily_team_queue,
|
||||
daily_org_spend_update_queue=daily_org_queue,
|
||||
daily_end_user_spend_update_queue=daily_end_user_queue,
|
||||
daily_agent_spend_update_queue=daily_agent_queue,
|
||||
daily_tag_spend_update_queue=daily_tag_queue,
|
||||
)
|
||||
|
||||
# Should be called exactly once (pipeline)
|
||||
mock_redis_cache.async_rpush_pipeline.assert_called_once()
|
||||
|
||||
# Verify only 3 operations were included (empty ones skipped)
|
||||
call_args = mock_redis_cache.async_rpush_pipeline.call_args
|
||||
rpush_list = call_args.kwargs["rpush_list"]
|
||||
assert len(rpush_list) == 3
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_in_memory_spend_updates_all_empty_returns_early(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
"""
|
||||
When all queues are empty, pipeline should never be called.
|
||||
"""
|
||||
mock_redis_cache.async_rpush_pipeline = AsyncMock()
|
||||
|
||||
# All queues return empty
|
||||
empty_queue = AsyncMock()
|
||||
empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
empty_daily_queue = AsyncMock()
|
||||
empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(
|
||||
return_value={}
|
||||
)
|
||||
|
||||
await redis_update_buffer.store_in_memory_spend_updates_in_redis(
|
||||
spend_update_queue=empty_queue,
|
||||
daily_spend_update_queue=empty_daily_queue,
|
||||
daily_team_spend_update_queue=empty_daily_queue,
|
||||
daily_org_spend_update_queue=empty_daily_queue,
|
||||
daily_end_user_spend_update_queue=empty_daily_queue,
|
||||
daily_agent_spend_update_queue=empty_daily_queue,
|
||||
daily_tag_spend_update_queue=empty_daily_queue,
|
||||
)
|
||||
|
||||
mock_redis_cache.async_rpush_pipeline.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_transactions_from_redis_buffer_pipeline(
|
||||
redis_update_buffer, mock_redis_cache
|
||||
):
|
||||
"""
|
||||
Verify get_all_transactions_from_redis_buffer_pipeline correctly parses
|
||||
and aggregates results from async_lpop_pipeline.
|
||||
"""
|
||||
# Simulate pipeline results: slot 0 = spend updates, slots 1-6 = daily categories
|
||||
db_spend_json = json.dumps(
|
||||
{
|
||||
"key_list_transactions": {"key1": 1.0, "key2": 2.0},
|
||||
"user_list_transactions": {"user1": 0.5},
|
||||
"end_user_list_transactions": {},
|
||||
"team_list_transactions": {},
|
||||
"team_member_list_transactions": {},
|
||||
"org_list_transactions": {},
|
||||
"tag_list_transactions": {},
|
||||
}
|
||||
)
|
||||
daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}})
|
||||
daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}})
|
||||
|
||||
mock_redis_cache.async_lpop_pipeline = AsyncMock(
|
||||
return_value=[
|
||||
[db_spend_json], # slot 0: db spend updates
|
||||
[daily_user_json], # slot 1: daily user
|
||||
[daily_team_json], # slot 2: daily team
|
||||
None, # slot 3: daily org (empty)
|
||||
None, # slot 4: daily end-user (empty)
|
||||
None, # slot 5: daily agent (empty)
|
||||
None, # slot 6: daily tag (empty)
|
||||
]
|
||||
)
|
||||
|
||||
result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
|
||||
assert len(result) == 7
|
||||
db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent, daily_tag = result
|
||||
|
||||
# Verify db spend was parsed correctly
|
||||
assert db_spend is not None
|
||||
assert db_spend["key_list_transactions"]["key1"] == 1.0
|
||||
assert db_spend["key_list_transactions"]["key2"] == 2.0
|
||||
assert db_spend["user_list_transactions"]["user1"] == 0.5
|
||||
|
||||
# Verify daily user was parsed
|
||||
assert daily_user is not None
|
||||
assert daily_user["user_key1"]["spend"] == 1.0
|
||||
|
||||
# Verify daily team was parsed
|
||||
assert daily_team is not None
|
||||
assert daily_team["team_key1"]["spend"] == 2.0
|
||||
|
||||
# Verify empty slots
|
||||
assert daily_org is None
|
||||
assert daily_end_user is None
|
||||
assert daily_agent is None
|
||||
assert daily_tag is None
|
||||
|
||||
# Verify pipeline was called once with correct keys
|
||||
mock_redis_cache.async_lpop_pipeline.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis():
|
||||
"""When redis_cache is None, should return all Nones"""
|
||||
buffer = RedisUpdateBuffer(redis_cache=None)
|
||||
result = await buffer.get_all_transactions_from_redis_buffer_pipeline()
|
||||
assert result == (None, None, None, None, None, None, None)
|
||||
@@ -1076,3 +1076,46 @@ async def test_commit_key_spend_updates_includes_last_active():
|
||||
last_active = call_kwargs["data"]["last_active"]
|
||||
assert isinstance(last_active, datetime)
|
||||
assert before_call <= last_active <= after_call
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_commit_spend_updates_uses_pipeline():
|
||||
"""
|
||||
Verify that _commit_spend_updates_to_db_with_redis uses
|
||||
get_all_transactions_from_redis_buffer_pipeline instead of 7 individual calls.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
mock_redis_update_buffer = AsyncMock()
|
||||
mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock()
|
||||
# Return all-None tuple (no data to commit)
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock(
|
||||
return_value=(None, None, None, None, None, None, None)
|
||||
)
|
||||
db_writer.redis_update_buffer = mock_redis_update_buffer
|
||||
|
||||
mock_pod_lock_manager = AsyncMock()
|
||||
mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True)
|
||||
mock_pod_lock_manager.release_lock = AsyncMock()
|
||||
db_writer.pod_lock_manager = mock_pod_lock_manager
|
||||
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_proxy_logging = MagicMock()
|
||||
|
||||
await db_writer._commit_spend_updates_to_db_with_redis(
|
||||
prisma_client=mock_prisma_client,
|
||||
n_retry_times=1,
|
||||
proxy_logging_obj=mock_proxy_logging,
|
||||
)
|
||||
|
||||
# Pipeline method should be called once
|
||||
mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_called_once()
|
||||
|
||||
# Individual methods should NOT be called
|
||||
mock_redis_update_buffer.get_all_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called()
|
||||
|
||||
Reference in New Issue
Block a user