diff --git a/tests/integration/observability/test_anthropic_overloaded_failure_signals.py b/tests/integration/observability/test_anthropic_overloaded_failure_signals.py index 85f452c4a19..3d6414dadd5 100644 --- a/tests/integration/observability/test_anthropic_overloaded_failure_signals.py +++ b/tests/integration/observability/test_anthropic_overloaded_failure_signals.py @@ -338,15 +338,15 @@ def test_the_datadog_proxy_failure_log_carries_the_answered_status_and_the_llm_f def _datadog_llm_failures(rig: _Rig, aliases: Sequence[str]) -> Mapping[str, Mapping[str, JsonValue]]: batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + def alias_of(message: Mapping[str, JsonValue]) -> str | None: + if "user_api_key_dict" in message: + return None + return next((alias for alias in aliases if _keyed_by(message, alias)), None) + def delivered() -> Mapping[str, Mapping[str, JsonValue]]: batches.extend(rig.datadog.drain()) - return { - alias: object_value(message["error_information"]) - for message in _datadog_messages(batches) - if "user_api_key_dict" not in message - for alias in aliases - if _keyed_by(message, alias) - } + keyed: Final = ((alias_of(message), message) for message in _datadog_messages(batches)) + return {alias: object_value(message["error_information"]) for alias, message in keyed if alias is not None} return eventually(delivered, lambda found: len(found) >= len(aliases), seconds=40, return_last_on_timeout=True) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index d729674288a..2073c61f2bb 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -9660,23 +9660,24 @@ class _RecordingLogging(Logging): messages=[{"role": "user", "content": "Hey"}], stream=True, call_type="completion", - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc), litellm_call_id="12345", function_id="1245", ) self.logged_failures: list[Exception] = [] # mutable-ok: test double accumulating callback arguments + self.failure_logged: Final = threading.Event() def failure_handler(self, exception: Exception, traceback_exception: str, *args: object, **kwargs: object) -> None: self.logged_failures.append(exception) + self.failure_logged.set() async def dispatch_attempt_failure_handlers(self, exception: Exception, traceback_exception: str) -> None: self.logged_failures.append(exception) + self.failure_logged.set() def _assert_failure_callback_got_the_public_exception(logging: _RecordingLogging, raised: MidStreamFallbackError) -> None: - deadline: Final = time.monotonic() + 5 - while not logging.logged_failures and time.monotonic() < deadline: - time.sleep(0.01) + assert logging.failure_logged.wait(timeout=5), "failure callbacks never ran" assert len(logging.logged_failures) == 1, logging.logged_failures logged: Final = logging.logged_failures[0] assert type(logged) is litellm.InternalServerError, logged