From 3bfeaa78ddd1bb7cae9ab6576e21a699a466b642 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:39:28 -0400 Subject: [PATCH] revert: undo accidental whole-file ruff-format of test_check_batch_cost.py The prior commit ran ruff format on the whole file to type the _batch_cost_result helper, reflowing hundreds of unrelated pre-existing lines that were never ruff-format-clean to begin with (confirmed at the PR's base commit, before any of these changes). CI's ruff-format gate only checks litellm/**/*.py, not tests/, so this reformatting served no CI purpose and only bloated the diff. Restores everything except the intended TYPE_CHECKING import and _batch_cost_result annotations. --- .../proxy_unit_tests/test_check_batch_cost.py | 507 ++++++++++++------ 1 file changed, 354 insertions(+), 153 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff2dce498f0..cd7f28007af 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -100,7 +100,9 @@ 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, ) @@ -112,15 +114,23 @@ 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,7 +139,9 @@ 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() @@ -140,7 +152,9 @@ 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") ) @@ -164,12 +178,18 @@ 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() @@ -195,7 +215,9 @@ 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"), []] @@ -203,7 +225,9 @@ 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 @@ -214,20 +238,32 @@ 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 @@ -241,9 +277,13 @@ 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" @@ -252,16 +292,22 @@ 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" @@ -309,7 +355,9 @@ 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() @@ -318,11 +366,15 @@ 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 @@ -398,9 +450,7 @@ 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", @@ -418,9 +468,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 @@ -434,9 +484,13 @@ 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" @@ -444,15 +498,21 @@ 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" @@ -500,7 +560,9 @@ 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() @@ -508,13 +570,15 @@ 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 @@ -658,16 +722,22 @@ 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" @@ -701,9 +771,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"]) @@ -720,9 +790,13 @@ 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" @@ -732,25 +806,31 @@ 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"]) @@ -785,9 +865,13 @@ 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 @@ -797,7 +881,9 @@ 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" @@ -806,7 +892,9 @@ 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", @@ -824,7 +912,9 @@ 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() @@ -875,9 +965,13 @@ 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" @@ -887,18 +981,24 @@ 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", @@ -906,18 +1006,22 @@ 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( @@ -928,7 +1032,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_job = MagicMock() @@ -936,7 +1042,9 @@ 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" @@ -962,9 +1070,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"]) @@ -981,9 +1089,13 @@ 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" @@ -991,15 +1103,21 @@ 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" @@ -1047,7 +1165,9 @@ 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() @@ -1055,16 +1175,20 @@ 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" - ) + 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 - 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 ( + 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" @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1081,9 +1205,13 @@ 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" @@ -1093,17 +1221,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( @@ -1124,10 +1258,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 @@ -1140,9 +1276,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=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" @@ -1151,7 +1291,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" @@ -1162,10 +1304,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" @@ -1180,7 +1326,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"}' @@ -1223,7 +1371,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() @@ -1274,7 +1424,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): @@ -1312,7 +1464,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): @@ -1332,7 +1486,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() @@ -1380,7 +1536,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, @@ -1388,7 +1546,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) @@ -1412,7 +1572,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"} @@ -1434,7 +1596,9 @@ 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) @@ -1465,7 +1629,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() @@ -1506,7 +1672,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): @@ -1561,7 +1729,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 @@ -1599,7 +1769,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 @@ -1607,7 +1779,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) @@ -1630,9 +1804,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" @@ -1648,7 +1826,9 @@ 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) @@ -1679,7 +1859,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() @@ -1803,7 +1985,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" @@ -1812,8 +1996,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 ) ) @@ -1882,7 +2066,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] @@ -1896,7 +2082,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] @@ -1904,7 +2092,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 @@ -1921,8 +2111,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.""" @@ -2018,7 +2206,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") @@ -2114,7 +2304,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() @@ -2152,7 +2344,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): @@ -2205,13 +2399,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): @@ -2266,11 +2464,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):