From 38145c20826697c13e3b1955e41d53061376eb94 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 13:17:43 -0700 Subject: [PATCH] test: undo the drive-by reformat below the poller error-file regression test --- .../proxy_unit_tests/test_check_batch_cost.py | 196 ++++++++++++------ 1 file changed, 137 insertions(+), 59 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 4f9fcf1952b..757762eac87 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -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: