headroom/tests/test_compression_summary_tool_eval.py
chopratejas d5ca50cd03 fix(tests): stop module-level dotenv loaders from polluting os.environ during pytest collection
# The bug

Several test modules and two production modules loaded the project `.env`
at *import time*. During pytest collection (where every test module is
imported once), this populated `os.environ` with API keys from `.env`.

The skipif guards in `test_proxy_passthrough_integration.py` (and
others) evaluate at collection time:

    @pytest.mark.skipif(not os.environ.get("OPENAI_API_KEY"), reason="...")

If the polluter module was collected *before* the guard, the guard saw
the leaked key, decided not to skip, and the integration tests ran
live against a fake key and failed. In a fresh local-dev venv with
`.env` + full `[dev]` extras, this manifested as ~16 spurious test
failures plus a misleading test runtime of 6+ minutes (live HTTP).

# Why now

CI does not see this (no `.env`). It only manifests when:
1. `litellm` (and friends) are installed — they run `dotenv.load_dotenv()`
   on import, populating `os.environ` from `.env`.
2. A `.env` file with real API keys exists locally.

Until the venv was provisioned with the full `[dev]` extras during
recent test work, `pytest.importorskip("litellm")` and
`from headroom.pricing import litellm_pricing` both silently no-op'd
(via try/except ImportError → `LITELLM_AVAILABLE=False`), so the leak
never triggered. With litellm now installed, the latent bug surfaced.

# The fix — three patterns

1. **Production modules** (`headroom/pricing/litellm_pricing.py`,
   `headroom/backends/litellm.py`): wrap the eager `import litellm` with
   a snapshot/restore of `os.environ`. Any keys litellm's bundled
   `python-dotenv` adds during import are deleted immediately. The
   module is fully imported and cached in `sys.modules` so subsequent
   imports hit the cache without re-running the side effect.

2. **Test modules using `pytest.importorskip("litellm")`**
   (`test_backend_bugs.py`, `test_bedrock_region.py`,
   `test_cost_tracker_counterfactual.py`): replace with
   `tests._dotenv.importorskip_no_env_leak("litellm")`, which does the
   same snapshot/restore around `importlib.import_module`.

3. **Test modules that intentionally need `.env` values for skipif
   guards** (`test_compression_summary_*.py`, `test_query_echo.py`,
   `test_cost_tracker_counterfactual.py`, `test_memory_usage_integration.py`,
   `test_bundled_tools_savings.py`): replace module-level
   `os.environ.setdefault(...)` / `dotenv.load_dotenv()` with
   `tests._dotenv.load_env_overrides()` (returns a local dict — does
   NOT mutate `os.environ`) plus `autouse_apply_env(...)` (function-
   scoped fixture that applies via `monkeypatch.setenv`, auto-cleaned
   at teardown). The skipif still works because
   `ANTHROPIC_KEY = os.environ.get(...) or _env_overrides.get(...)`
   reads from the local dict as fallback.

# Helper module

New `tests/_dotenv.py` exposes:
- `load_env_overrides() -> dict[str, str]` — read `.env` into a dict.
- `autouse_apply_env(overrides) -> fixture` — function-scoped autouse
  fixture that applies via `monkeypatch.setenv`.
- `importorskip_no_env_leak(module) -> module` — drop-in
  `pytest.importorskip` substitute that quarantines env mutations.

# Results

Local full-suite (excluding live-LLM and live-feed tests):
- Before: 46 failed, 4830 passed, 387s
- After:   2 failed, 4672 passed, 134s

The remaining 2 failures are unrelated environment-dependent tests
(missing `PIL` / Docker daemon).
2026-04-26 09:15:37 -07:00

289 lines
10 KiB
Python

"""Eval: Does the LLM invoke headroom_retrieve when summaries are present?
The REAL test — it's not enough for the LLM to know something is missing.
It must actually call the tool to fetch it.
Compares:
- WITH summary: LLM sees "2 failed, 1 error" → should call headroom_retrieve
- WITHOUT summary: LLM sees "[90 items compressed]" → likely does NOT call tool
Requires: ANTHROPIC_API_KEY in environment or .env file.
Run: python -m pytest tests/test_compression_summary_tool_eval.py -v -s
"""
from __future__ import annotations
import json
import os
import pytest
from tests._dotenv import autouse_apply_env, load_env_overrides
_env_overrides = load_env_overrides()
ANTHROPIC_KEY = os.environ.get("ANTHROPIC_API_KEY") or _env_overrides.get("ANTHROPIC_API_KEY", "")
apply_dotenv = autouse_apply_env(_env_overrides)
pytestmark = pytest.mark.skipif(
not ANTHROPIC_KEY,
reason="ANTHROPIC_API_KEY not set — skipping integration tests",
)
# The headroom_retrieve tool definition (same as what CCR injects)
HEADROOM_RETRIEVE_TOOL = {
"name": "headroom_retrieve",
"description": (
"Retrieve original uncompressed content from Headroom's compression cache. "
"Use this when you need more details from compressed data. "
"You can pass a query to search within the compressed content."
),
"input_schema": {
"type": "object",
"properties": {
"hash": {
"type": "string",
"description": "The hash key from the compression marker",
},
"query": {
"type": "string",
"description": "Optional search query to find specific items within the compressed data",
},
},
"required": ["hash"],
},
}
def _call_claude_with_tools(messages: list[dict], tools: list[dict], max_tokens: int = 300) -> dict:
"""Make a real Anthropic API call with tool use."""
import httpx
resp = httpx.post(
"https://api.anthropic.com/v1/messages",
headers={
"X-Api-Key": ANTHROPIC_KEY,
"anthropic-version": "2023-06-01",
"Content-Type": "application/json",
},
json={
"model": "claude-sonnet-4-5-20250929",
"max_tokens": max_tokens,
"messages": messages,
"tools": tools,
},
timeout=30,
)
return resp.json()
def _make_test_results(n: int = 100) -> list[dict]:
"""Test suite output with hidden failures in the compressed portion."""
results = []
for i in range(n):
result = {
"test_name": f"test_module_{i // 10}.test_case_{i}",
"status": "passed",
"duration_ms": 50 + i * 3,
}
if i == 42:
result["status"] = "failed"
result["error"] = "AssertionError: expected 200, got 401 in auth_middleware"
result["test_name"] = "test_auth.test_login_expired_token"
if i == 67:
result["status"] = "failed"
result["error"] = "TimeoutError: database pool exhausted after 30s"
result["test_name"] = "test_database.test_concurrent_connections"
if i == 88:
result["status"] = "error"
result["error"] = "ImportError: cannot import 'NewFeature'"
result["test_name"] = "test_features.test_new_feature_integration"
results.append(result)
return results
def _has_tool_use(response: dict) -> bool:
"""Check if the response contains a tool_use block."""
for block in response.get("content", []):
if block.get("type") == "tool_use":
return True
return False
def _get_tool_calls(response: dict) -> list[dict]:
"""Extract all tool_use blocks from response."""
calls = []
for block in response.get("content", []):
if block.get("type") == "tool_use":
calls.append(
{
"name": block.get("name"),
"input": block.get("input", {}),
}
)
return calls
class TestToolInvocationWithSummary:
"""The real eval: does the LLM call headroom_retrieve?"""
def test_with_summary_triggers_tool_call(self):
"""WITH compression summary → LLM should call headroom_retrieve."""
test_results = _make_test_results(100)
kept = test_results[:10] # All passing
from headroom.transforms.compression_summary import summarize_dropped_items
summary = summarize_dropped_items(test_results, kept)
compressed = json.dumps(kept, indent=2)
compressed += (
f"\n[90 items compressed to 10. Omitted: {summary}."
f' Retrieve specific items: headroom_retrieve(hash="ccr_test_abc123", query="your search")]'
)
messages = [
{
"role": "user",
"content": (
"Here are the test results from our CI pipeline:\n\n"
f"{compressed}\n\n"
"Tell me about any test failures. What went wrong?"
),
},
]
resp = _call_claude_with_tools(messages, [HEADROOM_RETRIEVE_TOOL])
tool_calls = _get_tool_calls(resp)
stop_reason = resp.get("stop_reason", "")
print(f"\n Summary: {summary}")
print(f" Stop reason: {stop_reason}")
print(f" Tool calls: {tool_calls}")
# With a summary showing failures, the LLM SHOULD call the tool
if stop_reason == "tool_use":
assert len(tool_calls) > 0
call = tool_calls[0]
assert call["name"] == "headroom_retrieve"
assert call["input"].get("hash") == "ccr_test_abc123"
# The query should be about failures/errors
query = call["input"].get("query", "").lower()
print(f" Query used: {query}")
has_relevant_query = any(
term in query for term in ["fail", "error", "issue", "problem", "broken", "test"]
)
assert has_relevant_query, f"Tool was called but query isn't relevant: {query}"
print(" RESULT: LLM invoked headroom_retrieve with relevant query ✓")
else:
# LLM responded with text — check if it at least mentions the failures
text = ""
for block in resp.get("content", []):
if block.get("type") == "text":
text += block.get("text", "")
print(f" LLM text response: {text[:200]}")
# It's acceptable if the LLM mentions it WANTS to retrieve
mentions_retrieval = any(
term in text.lower()
for term in ["retrieve", "headroom_retrieve", "fetch", "see more", "compressed"]
)
print(f" Mentions retrieval: {mentions_retrieval}")
def test_without_summary_baseline(self):
"""WITHOUT compression summary → LLM likely does NOT call tool."""
test_results = _make_test_results(100)
kept = test_results[:10] # All passing
compressed = json.dumps(kept, indent=2)
compressed += "\n[90 items compressed to 10. Retrieve more: hash=ccr_test_abc123]"
messages = [
{
"role": "user",
"content": (
"Here are the test results from our CI pipeline:\n\n"
f"{compressed}\n\n"
"Tell me about any test failures. What went wrong?"
),
},
]
resp = _call_claude_with_tools(messages, [HEADROOM_RETRIEVE_TOOL])
tool_calls = _get_tool_calls(resp)
stop_reason = resp.get("stop_reason", "")
print(f"\n Stop reason: {stop_reason}")
print(f" Tool calls: {tool_calls}")
if stop_reason == "tool_use":
call = tool_calls[0]
print(f" Query used: {call['input'].get('query', 'none')}")
print(" RESULT: LLM DID invoke tool (may check proactively)")
else:
text = ""
for block in resp.get("content", []):
if block.get("type") == "text":
text += block.get("text", "")
print(f" LLM text response: {text[:200]}")
print(" RESULT: LLM did NOT invoke tool — assumed all tests passed")
def test_code_summary_triggers_retrieval(self):
"""Code compression summary → LLM should retrieve specific function."""
compressed_code = '''class PaymentProcessor:
"""Processes payments via Stripe."""
def __init__(self, api_key: str):
# [2 lines omitted]
pass
def charge(self, amount: float, currency: str, token: str) -> dict:
# [8 lines omitted]
pass
def refund(self, charge_id: str, amount: float = None) -> dict:
# [3 lines omitted]
pass
def get_balance(self) -> float:
# [2 lines omitted]
pass
# [180 tokens compressed. removed: def charge (12 lines), def refund (6 lines). Retrieve full code: headroom_retrieve(hash="ccr_code_xyz", query="function name")]'''
messages = [
{
"role": "user",
"content": (
"Here's the payment processor code:\n\n"
f"```python\n{compressed_code}\n```\n\n"
"There's a bug in the retry logic for failed charges. "
"Can you find and fix it?"
),
},
]
resp = _call_claude_with_tools(messages, [HEADROOM_RETRIEVE_TOOL])
tool_calls = _get_tool_calls(resp)
stop_reason = resp.get("stop_reason", "")
print(f"\n Stop reason: {stop_reason}")
print(f" Tool calls: {tool_calls}")
if stop_reason == "tool_use":
call = tool_calls[0]
assert call["name"] == "headroom_retrieve"
query = call["input"].get("query", "").lower()
print(f" Query: {query}")
# Should be asking for the charge function specifically
has_charge = any(term in query for term in ["charge", "retry", "payment", "stripe"])
print(f" Targets charge/retry: {has_charge}")
print(" RESULT: LLM invoked tool to get the charge() implementation ✓")
else:
text = ""
for block in resp.get("content", []):
if block.get("type") == "text":
text += block.get("text", "")
print(f" LLM text: {text[:200]}")
print(" RESULT: LLM did NOT invoke tool")