diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 2d126fc3f29..97442f59a0d 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -123,16 +123,10 @@ class _RecordingLogger(CustomLogger): @pytest.fixture def recorder(): logger = _RecordingLogger() - saved_flag = ( - litellm.store_batch_line_items_in_callbacks - ) # test-quality-ok: process-wide opt-in flag; restored in teardown - saved_success = list( - litellm._async_success_callback - ) # test-quality-ok: the feature dispatches through this global list; restored in teardown + saved_flag = litellm.store_batch_line_items_in_callbacks # test-quality-ok: process-wide opt-in flag; restored in teardown + saved_success = list(litellm._async_success_callback) # test-quality-ok: the feature dispatches through this global list; restored in teardown saved_failure = list(litellm._async_failure_callback) # test-quality-ok: same dispatch seam, restored in teardown - litellm._async_success_callback = [ - logger - ] # test-quality-ok: there is no injection seam for callback lists; teardown restores + litellm._async_success_callback = [logger] # test-quality-ok: there is no injection seam for callback lists; teardown restores litellm._async_failure_callback = [logger] # test-quality-ok: same dispatch seam, restored in teardown yield logger litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: teardown restoring the value set above @@ -181,17 +175,11 @@ def _hidden(event: dict) -> dict: @pytest.mark.asyncio async def test_line_items_emitted_alongside_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it batch = _batch() with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(_parent_logging(), batch) @@ -227,9 +215,7 @@ async def test_line_items_emitted_alongside_aggregate(recorder): async def test_flag_off_emits_only_aggregate(recorder): assert litellm.store_batch_line_items_in_callbacks is False file_mock = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 @@ -239,13 +225,9 @@ async def test_flag_off_emits_only_aggregate(recorder): @pytest.mark.asyncio async def test_line_items_skipped_for_unsupported_provider(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(custom_llm_provider="cohere"), _batch()) file_mock.assert_not_awaited() @@ -257,9 +239,7 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): @pytest.mark.asyncio async def test_in_progress_batch_poll_emits_no_line_events(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it in_progress: Final = LiteLLMBatch( id="batch_wip", object="batch", @@ -272,9 +252,7 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): created_at=1, ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _parent_logging().async_success_handler(result=in_progress) file_mock.assert_not_called() @@ -284,19 +262,14 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) - with patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 assert _payload(recorder.success_events[0])["response_cost"] == 1.5 - EDGE_INPUT_JSONL = b"\n".join( [ json.dumps( @@ -533,9 +506,7 @@ def _parent_logging_with_params(litellm_params: dict) -> Logging: @pytest.mark.asyncio async def test_line_items_edge_shapes_and_edge_cases(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it batch = LiteLLMBatch( id="batch_edge", object="batch", @@ -553,16 +524,10 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, } ) - parent._litellm_internal_model_credentials = { - "api_key": "sk-line-items-marker" - } # test-quality-ok: private transport attribute, same channel the batch cost tracker uses + parent._litellm_internal_model_credentials = {"api_key": "sk-line-items-marker"} # test-quality-ok: private transport attribute, same channel the batch cost tracker uses with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) @@ -597,9 +562,7 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): @pytest.mark.asyncio async def test_line_items_child_params_drop_parent_credentials(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_file_content) parent: Final = _parent_logging_with_params( { @@ -611,17 +574,15 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): ) batch: Final = _batch() with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) line_events = [ - e for e in [*recorder.success_events, *recorder.failure_events] if _hidden(e).get("batch_custom_id") is not None + e + for e in [*recorder.success_events, *recorder.failure_events] + if _hidden(e).get("batch_custom_id") is not None ] assert len(line_events) == 2 for event in line_events: @@ -634,9 +595,7 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): @pytest.mark.asyncio async def test_line_items_anthropic_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it batch = LiteLLMBatch( id="batch_anth", object="batch", @@ -656,12 +615,8 @@ async def test_line_items_anthropic_shapes(recorder): custom_llm_provider="anthropic", ) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(logging_obj, batch) @@ -690,17 +645,11 @@ def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> L @pytest.mark.asyncio async def test_line_items_bedrock_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="bedrock"), @@ -737,17 +686,11 @@ async def test_line_items_bedrock_shapes(recorder): @pytest.mark.asyncio async def test_line_items_mistral_shape(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="mistral"), @@ -800,26 +743,24 @@ VERTEX_NATIVE_OUTPUT_JSONL = b"\n".join( ) +def _scoped_file_content(file_map): + return lambda file_id, **kwargs: SimpleNamespace(content=file_map[file_id]) + + @pytest.mark.asyncio async def test_line_items_anthropic_failure_keeps_provider_error(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock( - side_effect=lambda file_id, **kwargs: SimpleNamespace( - content={ + side_effect=_scoped_file_content( + { "input-anth-err": ANTHROPIC_ERR_INPUT_JSONL, "output-anth-err": ANTHROPIC_ERR_OUTPUT_JSONL, - }[file_id] + } ) ) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="anthropic"), @@ -835,24 +776,18 @@ async def test_line_items_anthropic_failure_keeps_provider_error(recorder): @pytest.mark.asyncio async def test_line_items_native_vertex_rows_are_skipped(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock( - side_effect=lambda file_id, **kwargs: SimpleNamespace( - content={ + side_effect=_scoped_file_content( + { "input-vtx": b"", "output-vtx": VERTEX_NATIVE_OUTPUT_JSONL, - }[file_id] + } ) ) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="vertex_ai"),