Skip to content

[Fix] Update check_responses_cost tests for _expire_stale_rows #25299

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuneng-berri merged 1 commit into from
Apr 7, 2026
Merged

[Fix] Update check_responses_cost tests for _expire_stale_rows #25299

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 37 additions & 37 deletions tests/proxy_unit_tests/test_check_responses_cost.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,11 +47,15 @@ def check_responses_cost_instance(
CheckResponsesCost,
)

return CheckResponsesCost(
instance = CheckResponsesCost(
proxy_logging_obj=mock_proxy_logging_obj,
prisma_client=mock_prisma_client,
llm_router=mock_llm_router,
)
# Mock _expire_stale_rows (raw SQL) so _cleanup_stale_managed_objects
# succeeds without a real DB. Individual tests can override this.
instance._expire_stale_rows = AsyncMock(return_value=0)
return instance

def test_initialization(self, check_responses_cost_instance):
"""Test that CheckResponsesCost initializes correctly"""
Expand All @@ -67,9 +71,6 @@ async def test_check_responses_cost_no_jobs(
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[]
)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=0
)

await check_responses_cost_instance.check_responses_cost()

Expand All @@ -86,24 +87,20 @@ async def test_check_responses_cost_no_jobs(
async def test_cleanup_stale_managed_objects(
self, check_responses_cost_instance, mock_prisma_client
):
"""Stale rows (older than cutoff) are bulk-updated to stale_expired before polling."""
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=5
)
"""Stale rows are expired via _expire_stale_rows before polling."""
from litellm.constants import STALE_OBJECT_CLEANUP_BATCH_SIZE

check_responses_cost_instance._expire_stale_rows = AsyncMock(return_value=5)
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[]
)

await check_responses_cost_instance.check_responses_cost()

# The first update_many call should be the stale-row cleanup scoped to "response"
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
stale_call = calls[0]
assert stale_call[1]["data"] == {"status": "stale_expired"}
where = stale_call[1]["where"]
assert where["file_purpose"] == "response"
assert "stale_expired" in where["status"]["not_in"]
assert "created_at" in where
# _expire_stale_rows should have been called with a cutoff datetime and batch size
check_responses_cost_instance._expire_stale_rows.assert_called_once()
call_args = check_responses_cost_instance._expire_stale_rows.call_args
assert call_args[0][1] == STALE_OBJECT_CLEANUP_BATCH_SIZE
Comment on lines +100 to +103

@greptile-apps greptile-apps Bot Apr 7, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Cutoff datetime argument not validated

The test verifies that _expire_stale_rows is called with the correct batch_size (second positional arg), but never asserts anything about the cutoff argument (first positional arg). A regression where a non-datetime value (e.g. a plain string or None) is passed instead of a timezone-aware datetime would go completely undetected.

Consider also asserting the type and timezone-awareness of the cutoff:

Suggested change
# _expire_stale_rows should have been called with a cutoff datetime and batch size
check_responses_cost_instance._expire_stale_rows.assert_called_once()
call_args = check_responses_cost_instance._expire_stale_rows.call_args
assert call_args[0][1] == STALE_OBJECT_CLEANUP_BATCH_SIZE
# _expire_stale_rows should have been called with a cutoff datetime and batch size
check_responses_cost_instance._expire_stale_rows.assert_called_once()
call_args = check_responses_cost_instance._expire_stale_rows.call_args
assert isinstance(call_args[0][0], datetime)
assert call_args[0][0].tzinfo is not None # must be timezone-aware
assert call_args[0][1] == STALE_OBJECT_CLEANUP_BATCH_SIZE

Rule Used: What: Flag any modifications to existing tests and... (source)


@pytest.mark.asyncio
async def test_check_responses_cost_with_completed_response(
Expand Down Expand Up @@ -145,10 +142,10 @@ async def test_check_responses_cost_with_completed_response(

await check_responses_cost_instance.check_responses_cost()

# calls[0] = stale cleanup, calls[1] = job completion
# update_many should only contain the job completion call
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 2
completion_call = calls[1]
assert len(calls) == 1
completion_call = calls[0]
assert completion_call[1]["data"]["status"] == "completed"
assert completion_call[1]["where"]["id"]["in"] == ["job-123"]

Expand Down Expand Up @@ -188,10 +185,10 @@ async def test_check_responses_cost_with_failed_response(

await check_responses_cost_instance.check_responses_cost()

# calls[0] = stale cleanup, calls[1] = job completion
# update_many should only contain the job completion call
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 2
assert calls[1][1]["data"]["status"] == "completed"
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"

@pytest.mark.asyncio
async def test_check_responses_cost_with_cancelled_response(
Expand Down Expand Up @@ -229,10 +226,10 @@ async def test_check_responses_cost_with_cancelled_response(

await check_responses_cost_instance.check_responses_cost()

# calls[0] = stale cleanup, calls[1] = job completion
# update_many should only contain the job completion call
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 2
assert calls[1][1]["data"]["status"] == "completed"
assert len(calls) == 1
assert calls[0][1]["data"]["status"] == "completed"

@pytest.mark.asyncio
async def test_check_responses_cost_with_in_progress_response(
Expand Down Expand Up @@ -270,10 +267,11 @@ async def test_check_responses_cost_with_in_progress_response(

await check_responses_cost_instance.check_responses_cost()

# Only the stale-cleanup call should have fired — no completion update
# No job completion update_many — response is still in progress
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 1
assert calls[0][1]["data"] == {"status": "stale_expired"}
assert len(calls) == 0
# Stale cleanup still ran via _expire_stale_rows
check_responses_cost_instance._expire_stale_rows.assert_called_once()

@pytest.mark.asyncio
async def test_check_responses_cost_with_queued_response(
Expand Down Expand Up @@ -311,10 +309,11 @@ async def test_check_responses_cost_with_queued_response(

await check_responses_cost_instance.check_responses_cost()

# Only the stale-cleanup call should have fired — no completion update
# No job completion update_many — response is still queued
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 1
assert calls[0][1]["data"] == {"status": "stale_expired"}
assert len(calls) == 0
# Stale cleanup still ran via _expire_stale_rows
check_responses_cost_instance._expire_stale_rows.assert_called_once()

@pytest.mark.asyncio
async def test_check_responses_cost_with_exception(
Expand Down Expand Up @@ -345,10 +344,11 @@ async def test_check_responses_cost_with_exception(
# Should not raise, just skip the job
await check_responses_cost_instance.check_responses_cost()

# Only the stale-cleanup call should have fired — no completion update
# No job completion update_many — exception skipped the job
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 1
assert calls[0][1]["data"] == {"status": "stale_expired"}
assert len(calls) == 0
# Stale cleanup still ran via _expire_stale_rows
check_responses_cost_instance._expire_stale_rows.assert_called_once()

@pytest.mark.asyncio
async def test_check_responses_cost_multiple_jobs(
Expand Down Expand Up @@ -424,10 +424,10 @@ async def test_check_responses_cost_multiple_jobs(

await check_responses_cost_instance.check_responses_cost()

# calls[0] = stale cleanup, calls[1] = completion of 2 finished jobs
# update_many should only contain the job completion call
calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list
assert len(calls) == 2
completion_call = calls[1]
assert len(calls) == 1
completion_call = calls[0]
assert len(completion_call[1]["where"]["id"]["in"]) == 2
assert "job-1" in completion_call[1]["where"]["id"]["in"]
assert "job-3" in completion_call[1]["where"]["id"]["in"]
Expand Down
Loading