test(logging): type the new MLflow, OTel and failure-handler tests

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

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(

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))