Spaces:
Running
Running
File size: 2,598 Bytes
c461ff0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | """Regression tests for proxy hook integration."""
from __future__ import annotations
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
pytest.importorskip("fastapi")
from fastapi.testclient import TestClient
from headroom.hooks import CompressionHooks
from headroom.proxy.savings_tracker import HEADROOM_SAVINGS_PATH_ENV_VAR
from headroom.proxy.server import ProxyConfig, create_app
def test_anthropic_hooks_do_not_break_extract_user_query_lookup(tmp_path, monkeypatch):
"""Hooks-enabled Anthropic requests should still reach the compression pipeline."""
monkeypatch.setenv(
HEADROOM_SAVINGS_PATH_ENV_VAR,
str(tmp_path / "proxy_savings.json"),
)
config = ProxyConfig(
cache_enabled=False,
rate_limit_enabled=False,
log_requests=False,
hooks=CompressionHooks(),
)
app = create_app(config)
with TestClient(app) as client:
proxy = client.app.state.proxy
large_user_content = "hello " * 200
proxy.anthropic_pipeline.apply = MagicMock(
return_value=SimpleNamespace(
messages=[{"role": "user", "content": "compressed"}],
transforms_applied=["test_transform"],
timing={},
tokens_before=100,
tokens_after=40,
waste_signals=None,
)
)
proxy._retry_request = AsyncMock(
return_value=httpx.Response(
200,
json={
"id": "msg_test",
"type": "message",
"role": "assistant",
"content": [{"type": "text", "text": "ok"}],
"usage": {
"input_tokens": 40,
"output_tokens": 5,
"cache_read_input_tokens": 0,
"cache_creation_input_tokens": 0,
},
},
)
)
response = client.post(
"/v1/messages",
json={
"model": "claude-sonnet-4-6",
"max_tokens": 128,
"messages": [{"role": "user", "content": large_user_content}],
},
)
assert response.status_code == 200
assert proxy.anthropic_pipeline.apply.called
assert response.headers["x-headroom-tokens-after"] == "40"
assert int(response.headers["x-headroom-tokens-before"]) > 40
assert int(response.headers["x-headroom-tokens-saved"]) > 0
|