diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 4eb9c7a5dfa..057c9978879 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -165,7 +165,9 @@ def _classify_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) - reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + reasoning_tokens: Final = ( + usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + ) yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 139bf583e6a..ff2dce498f0 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -6,14 +6,24 @@ Vertex (raw gs:// input_file_id) and Bedrock (raw s3:// input_file_id, ARN unified_object_id) batches with no managed unified id. """ +from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch import pytest +if TYPE_CHECKING: + from litellm.batches.batch_utils import BatchCostUsageResult + _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" -def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0): +def _batch_cost_result( + cost: float, + usage: dict, + models: list[str], + successful_requests: int = 1, + failed_requests: int = 0, +) -> "BatchCostUsageResult": """Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns, for mocking it in tests that only care about cost/usage/models.""" from litellm.batches.batch_utils import BatchCostUsageResult @@ -90,9 +100,7 @@ class TestCheckBatchCost: return MagicMock() @pytest.fixture - def check_batch_cost_instance( - self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router - ): + def check_batch_cost_instance(self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router): from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -104,23 +112,15 @@ class TestCheckBatchCost: ) @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list - ) + 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"] @@ -129,9 +129,7 @@ class TestCheckBatchCost: assert "created_at" in where @pytest.mark.asyncio - async def test_startup_probe_confirms_batch_processed_support( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) await check_batch_cost_instance.confirm_batch_processed_support() @@ -142,9 +140,7 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_startup_probe_marks_column_absent( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_marks_column_absent(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( side_effect=Exception("column batch_processed does not exist") ) @@ -168,18 +164,12 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_find_many_uses_pagination_and_excludes_stale( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_find_many_uses_pagination_and_excludes_stale(self, check_batch_cost_instance, mock_prisma_client): """find_many is called with take, order, and all terminal statuses excluded.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() @@ -205,9 +195,7 @@ class TestCheckBatchCost: """Falls back to query without batch_processed when primary query raises.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=[Exception("column batch_processed does not exist"), []] @@ -215,9 +203,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list - ) + calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list assert len(calls) == 2 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where @@ -228,32 +214,20 @@ class TestCheckBatchCost: assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio - async def test_column_absence_cached_across_cycles( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() # Only one find_many call — the fallback directly, no primary query attempt - assert ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 - ) - fallback_where = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1][ - "where" - ] - ) + assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -267,13 +241,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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-fallback-1" @@ -282,22 +252,16 @@ class TestCheckBatchCost: # Simulate column already known absent (e.g. discovered on a previous cycle) check_batch_cost_instance._has_batch_processed_column = False - 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]) # Build a fake batch response whose status triggers the completion branch mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - 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 = "openai" @@ -345,9 +309,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", 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() @@ -356,15 +318,11 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() # The update must have been called — this is the core assertion. - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - "batch_processed" not in update_data - ), "update() must NOT include batch_processed when column is absent" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert "batch_processed" not in update_data, "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -440,7 +398,9 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + return_value=_batch_cost_result( + 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"] + ), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -458,9 +418,9 @@ class TestCheckBatchCost: passed_kwargs = mock_afile_content.await_args[1] snapshot = passed_kwargs.get("_litellm_internal_model_credentials") assert snapshot is not None, "cost poller must pass the trusted credential snapshot" - assert isinstance( - snapshot, MappingProxyType - ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert isinstance(snapshot, MappingProxyType), ( + "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + ) assert snapshot["s3_bucket_name"] == "configured-batch-bucket" @pytest.mark.asyncio @@ -474,13 +434,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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-primary-1" @@ -488,21 +444,15 @@ 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]) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - 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 = "openai" @@ -550,9 +500,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", 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() @@ -560,15 +508,13 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - update_data["batch_processed"] is True - ), "update() must include batch_processed=True when column is present" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert update_data["batch_processed"] is True, ( + "update() must include batch_processed=True when column is present" + ) assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -712,22 +658,16 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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-anthropic-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - 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]) mock_response = MagicMock() mock_response.status = "completed" @@ -761,9 +701,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a failed cost tracking attempt must not mark the job processed" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a failed cost tracking attempt must not mark the job processed" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) @@ -780,13 +720,9 @@ class TestCheckBatchCost: """ import base64 - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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-terminal-1" @@ -796,31 +732,25 @@ 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]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), f"Expected update() to be called exactly once for a {terminal_status} job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + f"Expected update() to be called exactly once for a {terminal_status} job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == terminal_status - assert ( - update_data["batch_processed"] is True - ), "terminal-status update() must set batch_processed=True so polling stops" + assert update_data["batch_processed"] is True, ( + "terminal-status update() must set batch_processed=True so polling stops" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) @@ -855,13 +785,9 @@ class TestCheckBatchCost: f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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) input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id @@ -871,9 +797,7 @@ class TestCheckBatchCost: return input_file_row return None - mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file - ) + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=find_managed_file) mock_job = MagicMock() mock_job.id = "job-terminal-mint-1" @@ -882,9 +806,7 @@ class TestCheckBatchCost: mock_job.team_id = "team-1" 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]) response = LiteLLMBatch( id="batch-456", @@ -902,9 +824,7 @@ class TestCheckBatchCost: mock_hook = MagicMock() mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_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 await check_batch_cost_instance.check_batch_cost() @@ -955,13 +875,9 @@ class TestCheckBatchCost: import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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-completed-no-output-1" @@ -971,24 +887,18 @@ 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]) mock_response = MagicMock() mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{completed_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{completed_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) # Billing reads credentials off the router; if it is touched we billed a batch # that has no output, which is the behaviour this test guards against. - 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( "litellm.files.main.afile_content", @@ -996,22 +906,18 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a completed batch with no output file must be marked processed exactly once" - 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 completed batch with no output file must be marked processed exactly once" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == completed_status - assert ( - update_data["batch_processed"] is True - ), "completed-without-output update() must set batch_processed=True so polling stops" - assert ( - mock_afile_content.await_count == 0 - ), "a batch with no output file must not be billed" - assert ( - mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 - ), "a batch with no output file must not enter the cost-tracking path" + assert update_data["batch_processed"] is True, ( + "completed-without-output update() must set batch_processed=True so polling stops" + ) + assert mock_afile_content.await_count == 0, "a batch with no output file must not be billed" + assert mock_llm_router.get_deployment_credentials_with_provider.call_count == 0, ( + "a batch with no output file must not enter the cost-tracking path" + ) @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( @@ -1022,9 +928,7 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_job = MagicMock() @@ -1032,9 +936,7 @@ class TestCheckBatchCost: mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - 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]) mock_response = MagicMock() mock_response.status = "in_progress" @@ -1060,9 +962,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a non-terminal batch must not be written back (would stop polling prematurely)" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a non-terminal batch must not be written back (would stop polling prematurely)" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) @@ -1079,13 +981,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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-terminal-with-output-1" @@ -1093,21 +991,15 @@ 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]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' 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 = "openai" @@ -1155,9 +1047,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", 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() @@ -1165,20 +1055,16 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_afile_content.await_count == 1 - ), f"{terminal_status} batch with an output file must fetch results and be billed" - mock_logging_obj.async_success_handler.assert_awaited_once() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + assert mock_afile_content.await_count == 1, ( + f"{terminal_status} batch with an output file must fetch results and be billed" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + mock_logging_obj.async_success_handler.assert_awaited_once() + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["batch_processed"] is True - assert ( - update_data["status"] == terminal_status - ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + assert update_data["status"] == terminal_status, ( + f"billed {terminal_status} batch must keep its real terminal status in the DB" + ) @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1195,13 +1081,9 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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" @@ -1211,23 +1093,17 @@ 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( @@ -1248,12 +1124,10 @@ 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 @@ -1266,13 +1140,9 @@ 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=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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" @@ -1281,9 +1151,7 @@ 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" @@ -1294,14 +1162,10 @@ 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" @@ -1316,9 +1180,7 @@ 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"}' @@ -1361,9 +1223,7 @@ 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() @@ -1414,9 +1274,7 @@ 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): @@ -1454,9 +1312,7 @@ 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): @@ -1476,9 +1332,7 @@ 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() @@ -1526,9 +1380,7 @@ 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, @@ -1536,9 +1388,7 @@ 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) @@ -1562,9 +1412,7 @@ 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"} @@ -1586,9 +1434,7 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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) @@ -1619,9 +1465,7 @@ 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() @@ -1662,9 +1506,7 @@ 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): @@ -1719,9 +1561,7 @@ 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 @@ -1759,9 +1599,7 @@ 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 @@ -1769,9 +1607,7 @@ 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) @@ -1794,13 +1630,9 @@ 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" @@ -1816,9 +1648,7 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) 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) @@ -1849,9 +1679,7 @@ 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() @@ -1975,9 +1803,7 @@ 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" @@ -1986,8 +1812,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 ) ) @@ -2056,9 +1882,7 @@ 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] @@ -2072,9 +1896,7 @@ 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] @@ -2082,9 +1904,7 @@ 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 @@ -2101,6 +1921,8 @@ 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.""" @@ -2196,9 +2018,7 @@ 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") @@ -2294,9 +2114,7 @@ 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() @@ -2334,9 +2152,7 @@ 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): @@ -2389,17 +2205,13 @@ 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): @@ -2454,14 +2266,11 @@ 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):