scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 122

test(responses): drive the poller tests through the injected router

Browse source 
The eight claim/release tests reached into `litellm.aget_responses` with
`patch`, which the test-quality gate flags and which couples them to an
import path rather than to the poller's own seam. Each job now carries a
LiteLLM-encoded provider id, so the read routes through the router the
fixture already injects and the assertions run against that mock.

Claude-Session: https://claude.ai/code/session_01RHAjRxNhXTpKHeGMZ1nDKi
This commit is contained in:
ryan-crabbe-berri 2026-09-09 18:31:10 -07:00 • committed by jesus
parent 492336a50b
commit 5aa1901249
1 changed files with 58 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

119
tests/proxy_unit_tests/test_check_responses_cost.py
View file

@ -42,6 +42,17 @@ def _release_calls(mock_prisma_client):
)
def _routed_response_id(provider_response_id):
"""A LiteLLM-encoded id names a deployment, which is what sends the poll's read through the router."""
from litellm.responses.utils import ResponsesAPIRequestUtils
return ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai",
model_id="deployment-xyz",
response_id=provider_response_id,
)
class TestCheckResponsesCost:
"""Test suite for CheckResponsesCost class"""
@ -804,7 +815,7 @@ class TestCheckResponsesCost:
spend log, so losing the claim has to skip the read entirely or the job is billed twice."""
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_claimed_elsewhere"
mock_job.model_object_id = "resp_test_claimed_elsewhere"
mock_job.model_object_id = _routed_response_id("resp_test_claimed_elsewhere")
mock_job.created_by = "test-user"
mock_job.id = "job-claimed-elsewhere"
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_claimed_elsewhere"}
@ -816,10 +827,8 @@ class TestCheckResponsesCost:
return_value=0
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_sdk_aget:
await check_responses_cost_instance.check_responses_cost()
await check_responses_cost_instance.check_responses_cost()
mock_sdk_aget.assert_not_awaited()
mock_llm_router.aget_responses.assert_not_awaited()
assert _completion_calls(mock_prisma_client) == []
assert _release_calls(mock_prisma_client) == []
@ -832,7 +841,7 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_claim_is_taken_back_from_a_pod_that_died_holding_it(
self, check_responses_cost_instance, mock_prisma_client
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""A pod that dies between claiming and billing releases nothing, and the row's status
never reaches terminal, so without a lease every later cycle re-selects it and loses.
@ -845,7 +854,7 @@ class TestCheckResponsesCost:
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_abandoned"
mock_job.model_object_id = "resp_test_abandoned"
mock_job.model_object_id = _routed_response_id("resp_test_abandoned")
mock_job.created_by = "test-user"
mock_job.id = "job-abandoned"
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_abandoned"}
@ -866,9 +875,9 @@ class TestCheckResponsesCost:
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses = AsyncMock(return_value=mock_response)
await check_responses_cost_instance.check_responses_cost()
claim_where = _claim_calls(mock_prisma_client)[0].kwargs["where"]
abandoned_arm = next(arm for arm in claim_where["OR"] if "updated_at" in arm)
@ -881,13 +890,13 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_claim_is_taken_before_the_billing_read_and_kept_on_a_terminal_status(
self, check_responses_cost_instance, mock_prisma_client
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""The read prices the job, so the claim has to be taken before it, and keeping the claim
afterwards is what stops a second pod reading and billing the same row again."""
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_ordering"
mock_job.model_object_id = "resp_test_ordering"
mock_job.model_object_id = _routed_response_id("resp_test_ordering")
mock_job.created_by = "test-user"
mock_job.id = "job-ordering"
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_ordering"}
@ -919,10 +928,9 @@ class TestCheckResponsesCost:
side_effect=record_update_many
)
with patch(
"litellm.aget_responses", new_callable=AsyncMock, side_effect=record_read
):
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses = AsyncMock(side_effect=record_read)
await check_responses_cost_instance.check_responses_cost()
assert len(writes_and_reads) == 3
assert writes_and_reads[0] == {"batch_processed": True}
@ -932,13 +940,13 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
@pytest.mark.parametrize("provider_status", ["queued", "in_progress"])
async def test_non_terminal_status_releases_the_claim(
self, check_responses_cost_instance, mock_prisma_client, provider_status
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router, provider_status
):
"""A response the provider has not finished yet has no spend to record, so its row must go
back to batch_processed=False; holding the claim retires it before it is ever billed."""
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_still_running"
mock_job.model_object_id = "resp_test_still_running"
mock_job.model_object_id = _routed_response_id("resp_test_still_running")
mock_job.created_by = "test-user"
mock_job.id = "job-still-running"
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_still_running"}
@ -959,9 +967,9 @@ class TestCheckResponsesCost:
usage=None,
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses = AsyncMock(return_value=mock_response)
await check_responses_cost_instance.check_responses_cost()
assert _completion_calls(mock_prisma_client) == []
release_calls = _release_calls(mock_prisma_client)
@ -973,13 +981,13 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_failed_provider_read_releases_the_claim(
self, check_responses_cost_instance, mock_prisma_client
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""A read that raised billed nothing, so the claim has to be handed back or the row is
retired unbilled and no later poll cycle ever retries it."""
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_read_error"
mock_job.model_object_id = "resp_test_read_error"
mock_job.model_object_id = _routed_response_id("resp_test_read_error")
mock_job.created_by = "test-user"
mock_job.id = "job-read-error"
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_read_error"}
@ -991,12 +999,9 @@ class TestCheckResponsesCost:
return_value=1
)
with patch(
"litellm.aget_responses",
new_callable=AsyncMock,
side_effect=Exception("Provider error"),
):
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses = AsyncMock(side_effect=Exception("Provider error"))
await check_responses_cost_instance.check_responses_cost()
assert _completion_calls(mock_prisma_client) == []
release_calls = _release_calls(mock_prisma_client)
@ -1008,20 +1013,20 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_a_job_claimed_elsewhere_does_not_block_the_next_job(
self, check_responses_cost_instance, mock_prisma_client
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""Losing one row to another pod must skip only that row: the rest of the poll page still
has to be read and billed in the same cycle."""
mock_job1 = MagicMock()
mock_job1.unified_object_id = "resp_test_first"
mock_job1.model_object_id = "resp_test_first"
mock_job1.model_object_id = _routed_response_id("resp_test_first")
mock_job1.created_by = "user1"
mock_job1.id = "job-first"
mock_job1.file_object = {"model": "gpt-5", "id": "resp_test_first"}
mock_job2 = MagicMock()
mock_job2.unified_object_id = "resp_test_second"
mock_job2.model_object_id = "resp_test_second"
mock_job2.model_object_id = _routed_response_id("resp_test_second")
mock_job2.created_by = "user2"
mock_job2.id = "job-second"
mock_job2.file_object = {"model": "gpt-5", "id": "resp_test_second"}
@ -1042,12 +1047,14 @@ class TestCheckResponsesCost:
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses = AsyncMock(return_value=mock_response)
mock_aget.assert_awaited_once()
assert mock_aget.await_args.kwargs["response_id"] == "resp_test_second"
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses.assert_awaited_once()
assert mock_llm_router.aget_responses.await_args.kwargs["response_id"] == _routed_response_id(
"resp_test_second"
)
assert _completed_job_ids(mock_prisma_client) == ["job-second"]
@ -1095,13 +1102,13 @@ class TestCheckResponsesCost:
@pytest.mark.asyncio
async def test_old_schema_without_the_claim_column_still_bills_and_completes(
self, check_responses_cost_instance, mock_prisma_client
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""End to end on a pre-migration schema: the claim write fails, the response is still read
(which is what bills it) and the row is still marked completed."""
mock_job = MagicMock()
mock_job.unified_object_id = "resp_test_old_schema"
mock_job.model_object_id = "resp_test_old_schema"
mock_job.model_object_id = _routed_response_id("resp_test_old_schema")
mock_job.created_by = "test-user"
mock_job.id = "job-old-schema"
mock_job.file_object = {"model": "gpt-5", "id": "resp_test_old_schema"}
@ -1131,29 +1138,29 @@ class TestCheckResponsesCost:
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses = AsyncMock(return_value=mock_response)
mock_aget.assert_awaited_once()
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses.assert_awaited_once()
assert _completed_job_ids(mock_prisma_client) == ["job-old-schema"]
@pytest.mark.asyncio
async def test_a_failed_persist_does_not_abort_the_rest_of_the_poll_cycle(
self, check_responses_cost_instance, mock_prisma_client
self, check_responses_cost_instance, mock_prisma_client, mock_llm_router
):
"""One row's write failing must not take the whole cycle down with it: the jobs behind it
are already read and billed, so losing their write loses their usage for good."""
mock_job1 = MagicMock()
mock_job1.unified_object_id = "resp_test_persist_fails"
mock_job1.model_object_id = "resp_test_persist_fails"
mock_job1.model_object_id = _routed_response_id("resp_test_persist_fails")
mock_job1.created_by = "user1"
mock_job1.id = "job-persist-fails"
mock_job1.file_object = {"model": "gpt-5", "id": "resp_test_persist_fails"}
mock_job2 = MagicMock()
mock_job2.unified_object_id = "resp_test_persist_works"
mock_job2.model_object_id = "resp_test_persist_works"
mock_job2.model_object_id = _routed_response_id("resp_test_persist_works")
mock_job2.created_by = "user2"
mock_job2.id = "job-persist-works"
mock_job2.file_object = {"model": "gpt-5", "id": "resp_test_persist_works"}
@ -1174,11 +1181,11 @@ class TestCheckResponsesCost:
usage=ResponseAPIUsage(input_tokens=100, output_tokens=50, total_tokens=150),
)
with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget:
mock_aget.return_value = mock_response
await check_responses_cost_instance.check_responses_cost()
mock_llm_router.aget_responses = AsyncMock(return_value=mock_response)
assert mock_aget.await_count == 2
await check_responses_cost_instance.check_responses_cost()
assert mock_llm_router.aget_responses.await_count == 2
assert _completed_job_ids(mock_prisma_client) == [
"job-persist-fails",
"job-persist-works",
@ -1194,16 +1201,11 @@ class TestCheckResponsesCost:
is no handle on the generation.
"""
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.utils import SpecialEnums
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids")
provider_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai",
model_id="deployment-xyz",
response_id="resp_upstream_stable",
)
provider_response_id = _routed_response_id("resp_upstream_stable")
stale_advertised_id = "resp_" + str(
encrypt_value_helper(
value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(
@ -1243,16 +1245,11 @@ class TestCheckResponsesCost:
):
"""Rows created earlier carry the encrypted advertised id in both columns."""
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.utils import SpecialEnums
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key-for-response-ids")
provider_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai",
model_id="deployment-xyz",
response_id="resp_legacy_upstream",
)
provider_response_id = _routed_response_id("resp_legacy_upstream")
legacy_id = "resp_" + str(
encrypt_value_helper(
value=SpecialEnums.LITELLM_MANAGED_RESPONSE_API_RESPONSE_ID_COMPLETE_STR.value.format(