From 3aeca220319125a8616ed1236f4235de84082a38 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Fri, 13 Mar 2026 16:18:15 -0700
Subject: [PATCH] fix(test): update test_responses_background_cost assertions
 for pagination and stale cleanup

Tests were outdated after #23472 added pagination (take/order) to find_many
and stale-row cleanup via update_many. Updated assertions to match new call
signatures.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../test_responses_background_cost.py         | 50 ++++++++++++++-----
 1 file changed, 37 insertions(+), 13 deletions(-)

diff --git a/tests/test_litellm/integrations/test_responses_background_cost.py b/tests/test_litellm/integrations/test_responses_background_cost.py
index 4c4e9f36b26e..5cb42704181f 100644
--- a/tests/test_litellm/integrations/test_responses_background_cost.py
+++ b/tests/test_litellm/integrations/test_responses_background_cost.py
@@ -9,6 +9,7 @@
 
 import pytest
 
+from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE
 from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
 
 
@@ -336,12 +337,14 @@ async def test_check_responses_cost_no_jobs(
         # Should not raise any errors
         await checker.check_responses_cost()
 
-        # Verify find_many was called with correct parameters
+        # Verify find_many was called with correct parameters (includes pagination)
         mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with(
             where={
                 "status": {"in": ["queued", "in_progress"]},
                 "file_purpose": "response",
-            }
+            },
+            take=MAX_OBJECTS_PER_POLL_CYCLE,
+            order={"created_at": "asc"},
         )
 
     @pytest.mark.asyncio
@@ -394,12 +397,15 @@ async def test_check_responses_cost_with_completed_job(
             await checker.check_responses_cost()
 
             # Verify update_many was called to mark job as completed
-            mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
-            call_args = (
-                mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args
-            )
-            assert call_args[1]["where"]["id"]["in"] == ["job-123"]
-            assert call_args[1]["data"]["status"] == "completed"
+            # (stale cleanup also calls update_many, so check the specific completion call)
+            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            completion_calls = [
+                c for c in update_many_calls
+                if c.kwargs.get("where", {}).get("id") is not None
+            ]
+            assert len(completion_calls) == 1
+            assert completion_calls[0].kwargs["where"]["id"]["in"] == ["job-123"]
+            assert completion_calls[0].kwargs["data"]["status"] == "completed"
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_failed_job(
@@ -443,7 +449,13 @@ async def test_check_responses_cost_with_failed_job(
             await checker.check_responses_cost()
 
             # Verify job was marked as completed even though it failed
-            mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once()
+            # (stale cleanup also calls update_many, so check the specific completion call)
+            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            completion_calls = [
+                c for c in update_many_calls
+                if c.kwargs.get("where", {}).get("id") is not None
+            ]
+            assert len(completion_calls) == 1
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_with_in_progress_job(
@@ -486,8 +498,14 @@ async def test_check_responses_cost_with_in_progress_job(
 
             await checker.check_responses_cost()
 
-            # Verify update_many was NOT called (job still in progress)
-            mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
+            # Verify no completion update_many was called (job still in progress)
+            # (stale cleanup may still call update_many, so filter for completion calls)
+            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            completion_calls = [
+                c for c in update_many_calls
+                if c.kwargs.get("where", {}).get("id") is not None
+            ]
+            assert len(completion_calls) == 0
 
     @pytest.mark.asyncio
     async def test_check_responses_cost_error_handling(
@@ -524,5 +542,11 @@ async def test_check_responses_cost_error_handling(
             # Should not raise - errors are caught and logged
             await checker.check_responses_cost()
 
-            # Verify update_many was NOT called (error occurred)
-            mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()
+            # Verify no completion update_many was called (error occurred)
+            # (stale cleanup may still call update_many, so filter for completion calls)
+            update_many_calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
+            completion_calls = [
+                c for c in update_many_calls
+                if c.kwargs.get("where", {}).get("id") is not None
+            ]
+            assert len(completion_calls) == 0