chopratejas--headroom/tests/test_litellm_callback.py

"""Tests for HeadroomCallback LiteLLM integration.

Regression for #1114: HeadroomCallback did not inherit CustomLogger, so any
hook LiteLLM added post-1.89.x (e.g. async_post_call_success_hook) raised
AttributeError and crashed the LiteLLM proxy.
"""

from __future__ import annotations

import asyncio

from tests._dotenv import importorskip_no_env_leak

importorskip_no_env_leak("litellm")

from headroom.integrations.litellm_callback import HeadroomCallback  # noqa: E402


class TestHeadroomCallbackCustomLoggerInheritance:
    def test_instantiates_without_error(self) -> None:
        cb = HeadroomCallback()
        assert cb is not None

    def test_has_async_post_call_success_hook(self) -> None:
        """Regression: AttributeError: 'HeadroomCallback' has no attr 'async_post_call_success_hook'."""
        cb = HeadroomCallback()
        assert hasattr(cb, "async_post_call_success_hook"), (
            "async_post_call_success_hook must exist (added in litellm 1.89.x)"
        )

    def test_async_post_call_success_hook_is_callable(self) -> None:
        """LiteLLM must be able to await the hook without exception."""
        cb = HeadroomCallback()
        hook = cb.async_post_call_success_hook
        assert callable(hook)

    def test_async_post_call_success_hook_does_not_raise(self) -> None:
        """Calling the hook (no-op from CustomLogger) must not raise."""
        cb = HeadroomCallback()

        async def _run() -> None:
            await cb.async_post_call_success_hook(
                data={"model": "gpt-4o", "messages": []},
                user_api_key_dict={},
                response=None,
            )

        asyncio.run(_run())

    def test_all_current_litellm_async_hooks_present(self) -> None:
        """HeadroomCallback must expose every async hook CustomLogger defines."""
        from litellm.integrations.custom_logger import CustomLogger

        cb = HeadroomCallback()
        missing = [
            name
            for name in dir(CustomLogger)
            if name.startswith("async_") and not hasattr(cb, name)
        ]
        assert not missing, f"Missing CustomLogger hooks: {missing}"

    def test_async_pre_call_hook_still_works(self) -> None:
        """Inheritance must not break the existing compression hook."""
        cb = HeadroomCallback()
        assert hasattr(cb, "async_pre_call_hook")
        assert callable(cb.async_pre_call_hook)

    def test_total_tokens_saved_property(self) -> None:
        cb = HeadroomCallback()
        assert cb.total_tokens_saved == 0