From ec77cc1a1235e9c587f46f2d0f124d9e8819b600 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 5 Nov 2026 01:07:02 -0800 Subject: [PATCH] test(logging): type the new MLflow, OTel and failure-handler tests --- tests/unit/integrations/test_mlflow.py | 23 +++++++-------- tests/unit/integrations/test_opentelemetry.py | 10 +++---- .../test_litellm_logging.py | 28 +++++++++++++------ 3 files changed, 37 insertions(+), 24 deletions(-) diff --git a/tests/unit/integrations/test_mlflow.py b/tests/unit/integrations/test_mlflow.py index 516c24b27cc..0d86ef445a8 100644 --- a/tests/unit/integrations/test_mlflow.py +++ b/tests/unit/integrations/test_mlflow.py @@ -1,6 +1,7 @@ import asyncio import json import threading +from collections.abc import Mapping from datetime import datetime from typing import Final from unittest.mock import MagicMock, patch @@ -272,12 +273,12 @@ def test_mlflow_end_span_or_trace_works_with_mlflow_2x_client(): class _ExceptionSpanEvent: - def __init__(self, name, attributes): - self.name = name - self.attributes = attributes + def __init__(self, name: str, attributes: Mapping[str, str]) -> None: + self.name: Final = name + self.attributes: Final = attributes @classmethod - def from_exception(cls, exception): + def from_exception(cls, exception: BaseException) -> "_ExceptionSpanEvent": return cls( "exception", { @@ -289,18 +290,18 @@ class _ExceptionSpanEvent: @pytest.mark.parametrize("scrubbed", [True, False], ids=["scrubbed", "not_scrubbed"]) -def test_mlflow_failure_event_records_the_error_text_the_payload_logged(scrubbed): - secret = "my ssn is 123-45-6789" - modules = _mock_mlflow_modules() +def test_mlflow_failure_event_records_the_error_text_the_payload_logged(scrubbed: bool) -> None: + secret: Final = "my ssn is 123-45-6789" + modules: Final = _mock_mlflow_modules() modules["mlflow"].get_current_active_span.return_value = None modules["mlflow.entities"].SpanEvent = _ExceptionSpanEvent with patch.dict("sys.modules", modules): from litellm.integrations.mlflow import MlflowLogger - mlflow_logger = MlflowLogger() - span = MagicMock(parent_id=None, request_id="req-1") + mlflow_logger: Final = MlflowLogger() + span: Final = MagicMock(parent_id=None, request_id="req-1") mlflow_logger._client = MagicMock(**{"start_trace.return_value": span}) - error_message = "bad input: redacted-by-litellm" if scrubbed else f"bad input: {secret}" + error_message: Final = "bad input: redacted-by-litellm" if scrubbed else f"bad input: {secret}" mlflow_logger._handle_failure( kwargs={ @@ -317,7 +318,7 @@ def test_mlflow_failure_event_records_the_error_text_the_payload_logged(scrubbed end_time=datetime(2025, 1, 1), ) - event = span.add_event.call_args.args[0] + event: Final = span.add_event.call_args.args[0] mlflow_logger._client.end_trace.assert_called_once() assert event.attributes["exception.type"] == "ValueError" assert event.attributes["exception.message"] == error_message diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index e3a06acc716..2e8be029e9d 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -7638,18 +7638,18 @@ class TestOpentelemetryUnitTests: mock_span.set_attribute.assert_called_with(ErrorAttributes.ERROR_MESSAGE, "Fallback error message") -def test_record_exception_on_span_records_the_scrubbed_error_text_with_the_same_event_shape(): +def test_record_exception_on_span_records_the_scrubbed_error_text_with_the_same_event_shape() -> None: from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import SimpleSpanProcessor from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from litellm.integrations.opentelemetry import OpenTelemetry - exporter = InMemorySpanExporter() - provider = TracerProvider() + exporter: Final = InMemorySpanExporter() + provider: Final = TracerProvider() provider.add_span_processor(SimpleSpanProcessor(exporter)) - secret = "my ssn is 123-45-6789" - exception = ValueError(f"bad input: {secret}") + secret: Final = "my ssn is 123-45-6789" + exception: Final = ValueError(f"bad input: {secret}") with provider.get_tracer(__name__).start_as_current_span("failure") as span: OpenTelemetry()._record_exception_on_span( diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index cbbf27cde6d..0054982591d 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -12059,10 +12059,14 @@ class _FailureCapture(CustomLogger): self.sync_kwargs: dict | None = None self.async_kwargs: dict | None = None - def log_failure_event(self, kwargs, response_obj, start_time, end_time): + def log_failure_event( + self, kwargs: dict, response_obj: object, start_time: datetime.datetime, end_time: datetime.datetime + ) -> None: self.sync_kwargs = kwargs - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event( + self, kwargs: dict, response_obj: object, start_time: datetime.datetime, end_time: datetime.datetime + ) -> None: self.async_kwargs = kwargs @@ -12102,7 +12106,9 @@ async def _async_failure_payload( @pytest.mark.asyncio -async def test_failure_with_redaction_on_scrubs_the_prompt_out_of_the_error_and_keeps_the_rest(monkeypatch): +async def test_failure_with_redaction_on_scrubs_the_prompt_out_of_the_error_and_keeps_the_rest( + monkeypatch: pytest.MonkeyPatch, +) -> None: kwargs: Final = await _async_failure_payload(monkeypatch, redact=True, error=_prompt_echo_error()) payload: Final = kwargs["standard_logging_object"] @@ -12115,7 +12121,9 @@ async def test_failure_with_redaction_on_scrubs_the_prompt_out_of_the_error_and_ @pytest.mark.asyncio -async def test_failure_with_redaction_on_keeps_error_text_that_does_not_quote_the_prompt(monkeypatch): +async def test_failure_with_redaction_on_keeps_error_text_that_does_not_quote_the_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: error: Final = litellm.AuthenticationError( message="Incorrect API key provided", model="gpt-4o", llm_provider="openai" ) @@ -12128,7 +12136,7 @@ async def test_failure_with_redaction_on_keeps_error_text_that_does_not_quote_th @pytest.mark.asyncio -async def test_failure_with_redaction_off_logs_the_error_as_raised(monkeypatch): +async def test_failure_with_redaction_off_logs_the_error_as_raised(monkeypatch: pytest.MonkeyPatch) -> None: kwargs: Final = await _async_failure_payload(monkeypatch, redact=False, error=_prompt_echo_error()) assert SCRUB_PROMPT in kwargs["standard_logging_object"]["error_information"]["error_message"] @@ -12138,7 +12146,9 @@ async def test_failure_with_redaction_off_logs_the_error_as_raised(monkeypatch): @pytest.mark.asyncio @pytest.mark.parametrize("headers", [None, "x", ["a"]], ids=["none", "str", "list"]) @pytest.mark.parametrize("redact", [True, False], ids=["redaction_on", "redaction_off"]) -async def test_failure_with_non_mapping_headers_is_still_logged(monkeypatch, headers, redact): +async def test_failure_with_non_mapping_headers_is_still_logged( + monkeypatch: pytest.MonkeyPatch, headers: str | list[str] | None, redact: bool +) -> None: capture: Final = _FailureCapture() monkeypatch.setattr(litellm, "turn_off_message_logging", redact) monkeypatch.setattr(litellm, "_async_failure_callback", [capture]) @@ -12150,7 +12160,7 @@ async def test_failure_with_non_mapping_headers_is_still_logged(monkeypatch, hea assert (SCRUB_PROMPT in capture.async_kwargs["traceback_exception"]) is not redact -def test_sync_failure_hands_callbacks_the_scrubbed_error(monkeypatch): +def test_sync_failure_hands_callbacks_the_scrubbed_error(monkeypatch: pytest.MonkeyPatch) -> None: capture: Final = _FailureCapture() monkeypatch.setattr(litellm, "turn_off_message_logging", True) monkeypatch.setattr(litellm, "failure_callback", [capture]) @@ -12161,7 +12171,9 @@ def test_sync_failure_hands_callbacks_the_scrubbed_error(monkeypatch): assert SCRUB_PROMPT not in json.dumps(capture.sync_kwargs["standard_logging_object"], default=str) -def test_sync_failure_hands_lunary_the_scrubbed_traceback_when_only_the_traceback_quotes_the_prompt(monkeypatch): +def test_sync_failure_hands_lunary_the_scrubbed_traceback_when_only_the_traceback_quotes_the_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: tracked_events: Final[list[dict[str, object]]] = [] lunary_logger: Final = LunaryLogger() monkeypatch.setattr(lunary_logger.lunary_client, "track_event", lambda *_, **event: tracked_events.append(event))