[Fix] Update check_responses_cost tests for _expire_stale_rows

[yuneng-berri]

Summary

Problem

PR #25258 changed _cleanup_stale_managed_objects from using Prisma's update_many to raw SQL via a new _expire_stale_rows method, but the tests were not updated. This caused all 8 tests in test_check_responses_cost.py to fail in the auth-and-jwt CI job because:

• execute_raw was auto-generated as a MagicMock (not AsyncMock), causing await to fail
• The stale cleanup silently failed (caught by try/except), so update_many call counts didn't match expectations

Fix

• Mock _expire_stale_rows as AsyncMock on the fixture instance so _cleanup_stale_managed_objects succeeds without a real DB
• Update test_cleanup_stale_managed_objects to verify _expire_stale_rows is called with the correct batch size
• Adjust all update_many call count assertions: stale cleanup no longer goes through update_many, so terminal-state tests expect 1 call (not 2) and non-terminal tests expect 0 calls (not 1)

Testing

All 10 tests in test_check_responses_cost.py should now pass with the current _expire_stale_rows implementation on main.

Type

🐛 Bug Fix
✅ Test

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Apr 7, 2026 5:10pm

[codspeed-hq]

Merging this PR will not alter performance

✅ 16 untouched benchmarks

---

_{Comparing litellm_fix_check_responses_cost_tests (48a6823) with main (2bb7387)}

[greptile-apps]

Greptile Summary

This PR fixes 8 broken unit tests in test_check_responses_cost.py caused by PR #25258 which refactored _cleanup_stale_managed_objects from Prisma's update_many to raw SQL via a new _expire_stale_rows method. The tests were broken because execute_raw was auto-generated as a MagicMock (not AsyncMock), causing the await to fail silently inside a try/except, resulting in mismatched call-count assertions.

• Fixture now mocks _expire_stale_rows as AsyncMock(return_value=0) on the instance, ensuring stale cleanup succeeds without a real DB connection
• All update_many call-count assertions are reduced by 1 (terminal states: 2→1, non-terminal/exception: 1→0) since stale cleanup no longer routes through update_many
• test_cleanup_stale_managed_objects is updated to verify _expire_stale_rows is called exactly once with the correct STALE_OBJECT_CLEANUP_BATCH_SIZE
• Non-terminal-state and exception-path tests now explicitly assert _expire_stale_rows.assert_called_once() to confirm cleanup still runs even when no jobs complete
• Minor gap: test_cleanup_stale_managed_objects validates the batch_size arg to _expire_stale_rows but does not assert the cutoff argument is a timezone-aware datetime

Confidence Score: 5/5

Safe to merge — test-only fix that correctly aligns assertions with the actual implementation, with a single minor validation gap.

All changes are test-only and directly fix broken CI caused by a prior implementation change. Coverage is maintained or improved: stale-cleanup behavior is now tested via _expire_stale_rows assertions rather than indirectly through update_many counts. The only finding is a P2 style suggestion to also validate the cutoff datetime type, which does not block merge.

tests/proxy_unit_tests/test_check_responses_cost.py — minor: consider asserting isinstance(cutoff, datetime) and cutoff.tzinfo is not None in test_cleanup_stale_managed_objects

Important Files Changed

Filename	Overview
tests/proxy_unit_tests/test_check_responses_cost.py	Test suite correctly updated to reflect _expire_stale_rows replacing update_many for stale cleanup; call counts and assertions aligned with the new raw-SQL implementation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[check_responses_cost called] --> B[_cleanup_stale_managed_objects]
    B --> C[_expire_stale_rows\nraw SQL UPDATE via execute_raw]
    C -->|rows marked stale_expired| D[find_many: queued / in_progress jobs]
    D --> E{Any jobs?}
    E -->|No| Z[Done]
    E -->|Yes| F[Loop: aget_responses per job]
    F --> G{Terminal status?}
    G -->|completed / failed / cancelled| H[Add to completed_jobs]
    G -->|in_progress / queued| I[Skip job]
    H --> J{completed_jobs > 0?}
    J -->|Yes| K[update_many: status = completed]
    J -->|No| Z
    K --> Z

Loading

_{Reviews (1): Last reviewed commit: "fix(test): update check_responses_cost t..." | Re-trigger Greptile}

[greptile-apps]

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)