scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

test: undo the drive-by reformat below the poller error-file regression test

Browse source
This commit is contained in:
Mateo Wang 2026-08-29 13:17:43 -07:00
parent 99884f0eaa
commit 38145c2082
1 changed files with 137 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

196
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -1497,9 +1497,13 @@ class TestCheckBatchCost:
from litellm.exceptions import NotFoundError
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=1
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=None
)
mock_job = MagicMock()
mock_job.id = "job-output-gone-1"
@ -1509,17 +1513,23 @@ class TestCheckBatchCost:
mock_job.created_by = "user-1"
assert check_batch_cost_instance._has_batch_processed_column is True
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl"
mock_response = MagicMock()
mock_response.status = "failed"
mock_response.output_file_id = missing_output_file_id
mock_response.error_file_id = None
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}'
mock_response.model_dump_json.return_value = (
'{"id":"batch-1","status":"failed"}'
)
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(
return_value={"api_key": "sk-test"}
)
with (
patch(
@ -1540,10 +1550,12 @@ class TestCheckBatchCost:
assert mock_afile_content.await_count == 1
mock_calculate.assert_not_awaited()
assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, (
"a terminal batch with a 404ing output file must be retired, not retried forever"
)
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"]
assert (
mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
), "a terminal batch with a 404ing output file must be retired, not retried forever"
update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[
1
]["data"]
assert update_data["status"] == "failed"
assert update_data["batch_processed"] is True
@ -1556,9 +1568,13 @@ class TestCheckBatchCost:
Without this, GET /batches/{id} returns a raw file ID that cannot be routed
through the proxy, causing API_KEY errors when clients call GET /files/{id}/content.
"""
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
return_value=1
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=None
)
mock_job = MagicMock()
mock_job.id = "job-raw-file-1"
@ -1567,7 +1583,9 @@ class TestCheckBatchCost:
mock_job.team_id = None
check_batch_cost_instance._has_batch_processed_column = True
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job])
mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
raw_output_file_id = "file-batch-output-abc123"
raw_error_file_id = "file-batch-error-xyz456"
@ -1578,10 +1596,14 @@ class TestCheckBatchCost:
mock_response.status = "completed"
mock_response.output_file_id = raw_output_file_id
mock_response.error_file_id = raw_error_file_id
mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}'
mock_response.model_dump_json.return_value = (
'{"id":"batch-1","status":"completed"}'
)
mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
mock_llm_router.get_deployment_credentials_with_provider = MagicMock(
return_value={"api_key": "sk-test"}
)
mock_deployment = MagicMock()
mock_deployment.litellm_params.custom_llm_provider = "azure"
@ -1596,7 +1618,9 @@ class TestCheckBatchCost:
fake_managed_error_id,
]
mock_hook.store_unified_file_id = AsyncMock()
check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook
check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = (
mock_hook
)
mock_file_content = MagicMock()
mock_file_content.content = b'{"id":"req-1"}'
@ -1639,7 +1663,9 @@ class TestCheckBatchCost:
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("gpt-5-mini", "azure", None, None),
),
patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls,
patch(
"litellm.litellm_core_utils.litellm_logging.Logging"
) as mock_logging_cls,
):
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
@ -1690,7 +1716,9 @@ class TestUnmanagedVertexRouting:
def _job(self, file_object=None):
job = MagicMock()
job.unified_object_id = "8823717160934178816"
job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object()
job.file_object = (
file_object if file_object is not None else _unmanaged_vertex_file_object()
)
return job
def test_flag_off_skips_unmanaged_id_unchanged(self):
@ -1728,7 +1756,9 @@ class TestUnmanagedVertexRouting:
assert result == ("deploy-1", "8823717160934178816")
# bare model name (trailing GCS segment), not the full publishers/.. path
router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash")
router.resolve_model_name_from_model_id.assert_called_once_with(
"gemini-2.5-flash"
)
router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash")
def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self):
@ -1748,7 +1778,9 @@ class TestUnmanagedVertexRouting:
result = instance._resolve_job_routing(self._job(), prom)
assert result is None
prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment")
prom.record_check_batch_cost_error.assert_called_once_with(
"unmanaged_no_matching_deployment"
)
def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self):
router = MagicMock()
@ -1796,7 +1828,9 @@ class TestUnmanagedVertexRouting:
result = instance._resolve_job_routing(self._job(), prom)
assert result is None
prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment")
prom.record_check_batch_cost_error.assert_called_once_with(
"unmanaged_no_matching_deployment"
)
def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self):
"""Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable,
@ -1804,7 +1838,9 @@ class TestUnmanagedVertexRouting:
router = MagicMock()
instance = self._instance(track_unmanaged=True, router=router)
prom = MagicMock()
job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123"))
job = self._job(
file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")
)
with patch(_IS_B64, return_value=False):
result = instance._resolve_job_routing(job, prom)
@ -1828,7 +1864,9 @@ class TestUnmanagedVertexRouting:
mock_response.error_file_id = None
mock_response.completed_at = None
mock_response.created_at = None
mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}'
mock_response.model_dump_json.return_value = (
'{"id":"8823717160934178816","status":"completed"}'
)
router.aretrieve_batch = AsyncMock(return_value=mock_response)
router.get_deployment_credentials_with_provider = MagicMock(
return_value={"vertex_project": "p", "vertex_location": "us-central1"}
@ -1850,7 +1888,9 @@ class TestUnmanagedVertexRouting:
prisma.db.litellm_managedobjecttable = MagicMock()
prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
prisma.db.litellm_managedobjecttable.update = AsyncMock()
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()])
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[self._job()]
)
prisma.db.litellm_usertable = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
@ -1881,7 +1921,9 @@ class TestUnmanagedVertexRouting:
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("gemini-2.5-flash", "vertex_ai", None, None),
),
patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls,
patch(
"litellm.litellm_core_utils.litellm_logging.Logging"
) as mock_logging_cls,
):
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
@ -1922,7 +1964,9 @@ class TestUnmanagedBedrockRouting:
def _job(self, file_object=None):
job = MagicMock()
job.unified_object_id = self._ARN
job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object()
job.file_object = (
file_object if file_object is not None else _unmanaged_bedrock_file_object()
)
return job
def _bedrock_deployment(self):
@ -1977,7 +2021,9 @@ class TestUnmanagedBedrockRouting:
result = instance._resolve_job_routing(self._job(), prom)
assert result is None
prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment")
prom.record_check_batch_cost_error.assert_called_once_with(
"unmanaged_no_matching_deployment"
)
def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self):
"""The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured
@ -2015,7 +2061,9 @@ class TestUnmanagedBedrockRouting:
result = instance._resolve_job_routing(self._job(), prom)
assert result is None
prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment")
prom.record_check_batch_cost_error.assert_called_once_with(
"unmanaged_no_matching_deployment"
)
def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self):
"""Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as
@ -2023,7 +2071,9 @@ class TestUnmanagedBedrockRouting:
router = MagicMock()
instance = self._instance(track_unmanaged=True, router=router)
prom = MagicMock()
job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123"))
job = self._job(
file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")
)
with patch(_IS_B64, return_value=False):
result = instance._resolve_job_routing(job, prom)
@ -2046,9 +2096,13 @@ class TestUnmanagedBedrockRouting:
mock_response.error_file_id = None
mock_response.completed_at = None
mock_response.created_at = None
mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}'
mock_response.model_dump_json.return_value = (
f'{{"id":"{self._ARN}","status":"completed"}}'
)
router.aretrieve_batch = AsyncMock(return_value=mock_response)
router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"})
router.get_deployment_credentials_with_provider = MagicMock(
return_value={"aws_region_name": "us-east-1"}
)
deployment = self._bedrock_deployment()
deployment.model_name = "claude-sonnet-4"
@ -2064,7 +2118,9 @@ class TestUnmanagedBedrockRouting:
prisma.db.litellm_managedobjecttable = MagicMock()
prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1)
prisma.db.litellm_managedobjecttable.update = AsyncMock()
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()])
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[self._job()]
)
prisma.db.litellm_usertable = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
@ -2095,7 +2151,9 @@ class TestUnmanagedBedrockRouting:
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
return_value=("claude-sonnet-4", "bedrock", None, None),
),
patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls,
patch(
"litellm.litellm_core_utils.litellm_logging.Logging"
) as mock_logging_cls,
):
mock_logging_obj = MagicMock()
mock_logging_obj.async_success_handler = AsyncMock()
@ -2219,7 +2277,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
)
router = MagicMock()
router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
router.get_deployment_credentials_with_provider = MagicMock(
return_value={"api_key": "sk-test"}
)
deployment = MagicMock()
deployment.litellm_params.custom_llm_provider = "azure"
deployment.litellm_params.model = "azure/gpt-5.5"
@ -2228,8 +2288,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
router.get_deployment = MagicMock(return_value=deployment)
hook = MagicMock()
hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: (
_PROXY_LiteLLMManagedFiles.get_unified_output_file_id(
hook.get_unified_output_file_id = (
lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id(
None, output_file_id=output_file_id, model_id=model_id, model_name=model_name
)
)
@ -2298,7 +2358,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
get_models_from_unified_file_id,
)
output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)))
output_file_id = await self._run(
self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))
)
decoded = _is_base64_encoded_unified_file_id(output_file_id)
assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP]
@ -2312,7 +2374,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
_extract_models_from_managed_resource_id,
)
output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)))
output_file_id = await self._run(
self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))
)
models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None)
assert models == [self._PUBLIC_MODEL_GROUP]
@ -2320,7 +2384,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
await can_key_call_model(
model=models[0],
llm_model_list=None,
valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]),
valid_token=UserAPIKeyAuth(
api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]
),
llm_router=None,
)
is True
@ -2337,8 +2403,6 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
decoded = _is_base64_encoded_unified_file_id(output_file_id)
assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP]
class TestBatchCostAttribution:
"""CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so
the batch-cost log is attributed like a non-batch request."""
@ -2434,7 +2498,9 @@ class TestBatchCostAttribution:
"""An alias lookup failure must not lose the spend row; the key hash and team still
attribute it."""
instance = self._instance()
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down"))
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
side_effect=Exception("db down")
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
@ -2530,7 +2596,9 @@ class TestPollPageStarvation:
async def test_unified_id_without_model_id_is_retired(self):
"""A unified id that decodes but carries no model_id is unroutable no matter what
the config says, so it must leave the poll page instead of being retried forever."""
prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))])
prisma = self._prisma(
[self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]
)
llm_router = MagicMock()
llm_router.aretrieve_batch = AsyncMock()
@ -2568,7 +2636,9 @@ class TestPollPageStarvation:
await self._instance(prisma, llm_router).check_batch_cost()
prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True}
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {
"batch_processed": True
}
@pytest.mark.asyncio
async def test_provider_404_with_deployment_gone_keeps_job(self):
@ -2621,13 +2691,17 @@ class TestPollPageStarvation:
async def test_retirement_falls_back_to_status_without_batch_processed_column(self):
"""Older schemas have no batch_processed column, so the only way to stop selecting
the row is the status filter the poll query already applies."""
prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))])
prisma = self._prisma(
[self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]
)
instance = self._instance(prisma, MagicMock())
instance._has_batch_processed_column = False
await instance.check_batch_cost()
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"}
assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {
"status": "stale_expired"
}
@pytest.mark.asyncio
async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self):
@ -2682,11 +2756,14 @@ class TestPollPageStarvation:
await self._instance(prisma, llm_router).check_batch_cost()
retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list]
retired = [
call[1]["where"]["id"]
for call in prisma.db.litellm_managedobjecttable.update.call_args_list
]
assert retired == ["job-no-model", "job-gone"]
assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", (
"the newer healthy batch must still be polled in the same cycle"
)
assert (
llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live"
), "the newer healthy batch must still be polled in the same cycle"
@pytest.mark.asyncio
async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self):
@ -2715,7 +2792,6 @@ class TestPollPageStarvation:
prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
class _FakeManagedObjectRow:
"""One managed batch row the provider has finished but nothing has costed yet."""
@ -2732,12 +2808,8 @@ class _FakeManagedObjectRow:
self.request_tags = None
self.created_at = 1700000000
self.file_object = json.dumps(
{
"id": "batch-456",
"status": "in_progress",
"input_file_id": "file-input-1",
"output_file_id": _CLAIM_OUTPUT_FILE_ID,
}
{"id": "batch-456", "status": "in_progress", "input_file_id": "file-input-1",
"output_file_id": _CLAIM_OUTPUT_FILE_ID}
)
@ -2837,7 +2909,9 @@ class TestMultiPodBatchCostClaim:
router = MagicMock()
router.aretrieve_batch = AsyncMock(return_value=response)
router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"})
router.get_deployment_credentials_with_provider = MagicMock(
return_value={"api_key": "sk-test"}
)
router.get_deployment = MagicMock(return_value=deployment)
return router
@ -2997,7 +3071,9 @@ class TestMultiPodBatchCostClaim:
await asyncio.Event().wait()
with self._billing_patches(journal, during_fetch=_never_returns) as logging_obj:
interrupted = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost())
interrupted = asyncio.create_task(
self._instance(prisma, self._router()).check_batch_cost()
)
await asyncio.wait_for(reached_fetch.wait(), timeout=5)
assert row.batch_processed is False, "an in-flight costing must not mark the row processed"
interrupted.cancel()
@ -3032,7 +3108,9 @@ class TestMultiPodBatchCostClaim:
await finish_fetch.wait()
with self._billing_patches(journal, during_fetch=_wait_for_the_delete_attempt):
costing = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost())
costing = asyncio.create_task(
self._instance(prisma, self._router()).check_batch_cost()
)
await asyncio.wait_for(reached_fetch.wait(), timeout=5)
with pytest.raises(HTTPException) as blocked: