diff --git a/tests/test_litellm/integrations/focus/test_vantage_destination.py b/tests/test_litellm/integrations/focus/test_vantage_destination.py index 181d3d1381..c3fa9bda63 100644 --- a/tests/test_litellm/integrations/focus/test_vantage_destination.py +++ b/tests/test_litellm/integrations/focus/test_vantage_destination.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Any, Dict, List from unittest.mock import AsyncMock, patch @@ -12,12 +12,13 @@ from litellm.integrations.focus.destinations.base import FocusTimeWindow from litellm.integrations.focus.destinations.vantage_destination import ( FocusVantageDestination, VANTAGE_MAX_BYTES_PER_UPLOAD, + VANTAGE_MAX_ROWS_PER_UPLOAD, ) def _window(freq: str = "hourly", hour: int = 5) -> FocusTimeWindow: start = datetime(2024, 1, 2, hour, tzinfo=timezone.utc) - end = start.replace(hour=hour + 1) + end = start + timedelta(hours=1) return FocusTimeWindow(start_time=start, end_time=end, frequency=freq) @@ -136,3 +137,51 @@ async def test_should_batch_large_content(): # Each upload should be within limits for chunk in upload_calls: assert len(chunk) <= VANTAGE_MAX_BYTES_PER_UPLOAD + + +@pytest.mark.asyncio +async def test_should_batch_by_row_count(): + """Verify batching triggers when row count exceeds 10K even if under 2 MB.""" + dest = FocusVantageDestination(prefix="exports", config=_config()) + + header = b"col1" + # Short rows so total size stays well under 2 MB + row = b"x" + num_rows = VANTAGE_MAX_ROWS_PER_UPLOAD + 500 + content = header + b"\n" + b"\n".join([row] * num_rows) + b"\n" + + # Confirm content is under 2 MB but over 10K rows + assert len(content) < VANTAGE_MAX_BYTES_PER_UPLOAD + assert num_rows > VANTAGE_MAX_ROWS_PER_UPLOAD + + upload_calls: List[bytes] = [] + + mock_response = AsyncMock() + mock_response.raise_for_status = lambda: None + + mock_client = AsyncMock() + mock_client.__aenter__ = AsyncMock(return_value=mock_client) + mock_client.__aexit__ = AsyncMock(return_value=False) + + async def capture_post(url, **kwargs): + files = kwargs.get("files", {}) + if "file" in files: + upload_calls.append(files["file"][1]) + return mock_response + + mock_client.post = capture_post + + with patch("litellm.integrations.focus.destinations.vantage_destination.httpx.AsyncClient", return_value=mock_client): + await dest.deliver( + content=content, + time_window=_window(), + filename="usage.csv", + ) + + # Should have made at least 2 uploads due to row count + assert len(upload_calls) >= 2 + # Each batch should have at most 10K data rows (header + data rows + trailing newline) + for chunk in upload_calls: + lines = chunk.split(b"\n") + data_lines = [line for line in lines[1:] if line.strip()] + assert len(data_lines) <= VANTAGE_MAX_ROWS_PER_UPLOAD