chopratejas--headroom/tests/test_gemini_compression_offload.py

"""Gemini compression offload (perf): the 3 Gemini handlers must run the CPU-bound
`openai_pipeline.apply()` on the compression executor, not inline on the event loop.

The wiring (each handler awaits `_run_compression_in_executor(lambda: apply(...))`) mirrors
the proven openai/anthropic paths; these tests assert the two observable properties that
wiring delivers — apply runs on a worker thread, and the loop stays responsive during a
slow compression — plus a sanity check that the handlers are async and import the timeout.
"""

from __future__ import annotations

import asyncio
import inspect
import threading
import time

from headroom.proxy.server import ProxyConfig, create_app


def _make_proxy():  # noqa: ANN202 — returns the internal HeadroomProxy
    app = create_app(
        ProxyConfig(
            optimize=True,
            cache_enabled=False,
            rate_limit_enabled=False,
            cost_tracking_enabled=False,
        )
    )
    return app.state.proxy


def test_gemini_handlers_are_async_and_import_the_timeout() -> None:
    """Wiring sanity: the offload uses `await`, so the handlers must be coroutines, and the
    timeout constant must be importable in the module (a missing import would NameError)."""
    from headroom.proxy.handlers import gemini

    for name in (
        "handle_gemini_generate_content",
        "handle_google_cloudcode_stream",
        "handle_gemini_count_tokens",
    ):
        fn = getattr(gemini.GeminiHandlerMixin, name)
        assert inspect.iscoroutinefunction(fn), f"{name} must be async to await the offload"

    assert hasattr(gemini, "COMPRESSION_TIMEOUT_SECONDS")


async def test_compression_offload_runs_on_worker_thread() -> None:
    """apply() runs on a 'headroom-compress' executor thread, not the event-loop thread."""
    proxy = _make_proxy()
    loop_thread_name = threading.current_thread().name
    seen: dict[str, str] = {}

    def _slow_apply() -> str:
        seen["thread"] = threading.current_thread().name
        time.sleep(0.1)
        return "compressed"

    result = await proxy._run_compression_in_executor(_slow_apply, timeout=10)

    assert result == "compressed"
    assert seen["thread"].startswith("headroom-compress")
    assert seen["thread"] != loop_thread_name


async def test_compression_offload_keeps_event_loop_responsive() -> None:
    """While a slow compression runs on the executor, the loop keeps scheduling coroutines.
    A bare sync apply() on the loop (the bug this fixes) would starve them to ~0 ticks."""
    proxy = _make_proxy()
    ticks = 0

    async def _ticker() -> None:
        nonlocal ticks
        while True:
            await asyncio.sleep(0.01)
            ticks += 1

    def _slow_apply() -> str:
        time.sleep(0.3)
        return "x"

    tick_task = asyncio.create_task(_ticker())
    try:
        result = await proxy._run_compression_in_executor(_slow_apply, timeout=10)
    finally:
        tick_task.cancel()

    assert result == "x"
    # ~30 ticks expected at 10ms over 0.3s; a blocked loop would yield near zero.
    assert ticks >= 5