mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
Fix ruff lint errors in test files
This commit is contained in:
chopratejas
parent
e2e74a008a
commit
d9cc4f3991
5 changed files with 55 additions and 52 deletions
|
|
@ -20,7 +20,6 @@ from headroom.backends.litellm import (
|
|||
_normalize_bedrock_profile_id,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Region Prefix Mapping
|
||||
# =============================================================================
|
||||
|
|
@ -123,12 +122,12 @@ class TestFetchBedrockInferenceProfiles:
|
|||
"""Should return static map when boto3 is not installed."""
|
||||
with patch.dict("sys.modules", {"boto3": None}):
|
||||
# Force reimport failure
|
||||
import importlib
|
||||
|
||||
import headroom.backends.litellm as mod
|
||||
|
||||
# Temporarily break boto3 import inside the function
|
||||
original_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__
|
||||
original_import = (
|
||||
__builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__
|
||||
)
|
||||
|
||||
def mock_import(name, *args, **kwargs):
|
||||
if name == "boto3":
|
||||
raise ImportError("No module named 'boto3'")
|
||||
|
|
@ -155,12 +154,13 @@ class TestFetchBedrockInferenceProfiles:
|
|||
|
||||
with patch("headroom.backends.litellm.boto3", mock_boto3, create=True):
|
||||
# Patch the import inside the function
|
||||
original_fn = _fetch_bedrock_inference_profiles.__code__
|
||||
_fetch_bedrock_inference_profiles.__code__ # noqa: B018
|
||||
_bedrock_profiles_cache.clear()
|
||||
|
||||
# We need to actually test the function, so let's just use the
|
||||
# mock_boto3 and make sure the function catches the exception
|
||||
import builtins
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def patched_import(name, *args, **kwargs):
|
||||
|
|
@ -191,6 +191,7 @@ class TestFetchBedrockInferenceProfiles:
|
|||
mock_boto3.client.return_value = mock_client
|
||||
|
||||
import builtins
|
||||
|
||||
real_import = builtins.__import__
|
||||
|
||||
def patched_import(name, *args, **kwargs):
|
||||
|
|
@ -339,9 +340,10 @@ class TestNormalizeBedrockProfileId:
|
|||
)
|
||||
|
||||
def test_with_bedrock_slash_prefix(self):
|
||||
assert _normalize_bedrock_profile_id(
|
||||
"bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0"
|
||||
) == "claude-sonnet-4-20250514"
|
||||
assert (
|
||||
_normalize_bedrock_profile_id("bedrock/eu.anthropic.claude-sonnet-4-20250514-v1:0")
|
||||
== "claude-sonnet-4-20250514"
|
||||
)
|
||||
|
||||
def test_non_claude_returns_none(self):
|
||||
assert _normalize_bedrock_profile_id("eu.meta.llama-3-70b-v1:0") is None
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ import pytest
|
|||
|
||||
from headroom.memory.budget import (
|
||||
BudgetConfig,
|
||||
BudgetReport,
|
||||
MemoryBudgetManager,
|
||||
)
|
||||
from headroom.memory.writers.base import MemoryEntry
|
||||
|
|
@ -49,7 +48,10 @@ class TestBudgetManager:
|
|||
|
||||
def test_budget_limits(self, manager: MemoryBudgetManager):
|
||||
# Create many entries that exceed budget
|
||||
entries = [_make_entry(f"A long memory content string number {i} " * 10, importance=0.8) for i in range(50)]
|
||||
entries = [
|
||||
_make_entry(f"A long memory content string number {i} " * 10, importance=0.8)
|
||||
for i in range(50)
|
||||
]
|
||||
optimized, report = manager.optimize(entries, "claude") # 2000 token budget
|
||||
|
||||
assert report.pruned_budget > 0
|
||||
|
|
@ -77,8 +79,12 @@ class TestBudgetManager:
|
|||
|
||||
def test_merge_similar(self, manager: MemoryBudgetManager):
|
||||
entries = [
|
||||
_make_entry("Use source .venv/bin/activate && pytest for running tests", importance=0.5),
|
||||
_make_entry("Use source .venv/bin/activate && pytest for running tests", importance=0.8),
|
||||
_make_entry(
|
||||
"Use source .venv/bin/activate && pytest for running tests", importance=0.5
|
||||
),
|
||||
_make_entry(
|
||||
"Use source .venv/bin/activate && pytest for running tests", importance=0.8
|
||||
),
|
||||
_make_entry("Something completely different about architecture", importance=0.6),
|
||||
]
|
||||
optimized, report = manager.optimize(entries, "generic")
|
||||
|
|
@ -101,7 +107,10 @@ class TestBudgetManager:
|
|||
|
||||
def test_report_tokens(self, manager: MemoryBudgetManager):
|
||||
# Use entries large enough to trigger budget pruning
|
||||
entries = [_make_entry(f"A very long memory content entry {i} " * 30, importance=0.8) for i in range(20)]
|
||||
entries = [
|
||||
_make_entry(f"A very long memory content entry {i} " * 30, importance=0.8)
|
||||
for i in range(20)
|
||||
]
|
||||
_, report = manager.optimize(entries, "claude") # 2000 token budget
|
||||
|
||||
assert report.tokens_before > 0
|
||||
|
|
|
|||
|
|
@ -11,11 +11,8 @@ Verifies that:
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.proxy.server import HeadroomProxy, ProxyConfig
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# ProxyConfig Flag Resolution Tests
|
||||
# =============================================================================
|
||||
|
|
@ -217,13 +214,12 @@ class TestWrapCLILearnFlag:
|
|||
|
||||
def test_start_proxy_builds_learn_command(self):
|
||||
"""_start_proxy with learn=True adds --learn to command."""
|
||||
import sys
|
||||
|
||||
from headroom.cli.wrap import _start_proxy
|
||||
|
||||
# We can't actually run the proxy, but we can check the function signature
|
||||
import inspect
|
||||
|
||||
from headroom.cli.wrap import _start_proxy
|
||||
|
||||
sig = inspect.signature(_start_proxy)
|
||||
assert "learn" in sig.parameters
|
||||
|
||||
|
|
|
|||
|
|
@ -6,9 +6,6 @@ a real memory backend.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.memory.traffic_learner import (
|
||||
|
|
@ -19,7 +16,6 @@ from headroom.memory.traffic_learner import (
|
|||
_is_error,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Error Classification Tests
|
||||
# =============================================================================
|
||||
|
|
@ -124,9 +120,11 @@ class TestTrafficLearner:
|
|||
@pytest.mark.asyncio
|
||||
async def test_preference_extraction(self, learner: TrafficLearner):
|
||||
"""Test extraction of user preference signals."""
|
||||
await learner.on_messages([
|
||||
{"role": "user", "content": "don't use git push, I'll push manually"},
|
||||
])
|
||||
await learner.on_messages(
|
||||
[
|
||||
{"role": "user", "content": "don't use git push, I'll push manually"},
|
||||
]
|
||||
)
|
||||
|
||||
stats = learner.get_stats()
|
||||
assert stats["patterns_extracted"] >= 1
|
||||
|
|
@ -134,14 +132,16 @@ class TestTrafficLearner:
|
|||
@pytest.mark.asyncio
|
||||
async def test_preference_from_content_blocks(self, learner: TrafficLearner):
|
||||
"""Test preference extraction from Anthropic content block format."""
|
||||
await learner.on_messages([
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "stop running the full test suite without asking"},
|
||||
],
|
||||
},
|
||||
])
|
||||
await learner.on_messages(
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "stop running the full test suite without asking"},
|
||||
],
|
||||
},
|
||||
]
|
||||
)
|
||||
|
||||
stats = learner.get_stats()
|
||||
assert stats["patterns_extracted"] >= 1
|
||||
|
|
|
|||
|
|
@ -5,12 +5,9 @@ from __future__ import annotations
|
|||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.memory.writers.base import (
|
||||
MARKER_END,
|
||||
MARKER_START,
|
||||
ExportResult,
|
||||
MemoryEntry,
|
||||
_estimate_tokens,
|
||||
_merge_section,
|
||||
|
|
@ -20,7 +17,6 @@ from headroom.memory.writers.codex_writer import CodexMemoryWriter
|
|||
from headroom.memory.writers.cursor_writer import CursorMemoryWriter
|
||||
from headroom.memory.writers.generic_writer import GenericMemoryWriter
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Test Data
|
||||
# =============================================================================
|
||||
|
|
@ -31,14 +27,16 @@ def _make_entries(count: int = 5) -> list[MemoryEntry]:
|
|||
entries = []
|
||||
categories = ["error_recovery", "environment", "preference", "architecture"]
|
||||
for i in range(count):
|
||||
entries.append(MemoryEntry(
|
||||
content=f"Test memory entry {i}: use pytest not unittest",
|
||||
importance=0.5 + (i % 3) * 0.15,
|
||||
category=categories[i % len(categories)],
|
||||
entity_refs=[f"/path/to/file{i}.py"],
|
||||
created_at=time.time() - i * 3600, # Each an hour older
|
||||
access_count=max(0, 3 - i),
|
||||
))
|
||||
entries.append(
|
||||
MemoryEntry(
|
||||
content=f"Test memory entry {i}: use pytest not unittest",
|
||||
importance=0.5 + (i % 3) * 0.15,
|
||||
category=categories[i % len(categories)],
|
||||
entity_refs=[f"/path/to/file{i}.py"],
|
||||
created_at=time.time() - i * 3600, # Each an hour older
|
||||
access_count=max(0, 3 - i),
|
||||
)
|
||||
)
|
||||
return entries
|
||||
|
||||
|
||||
|
|
@ -94,9 +92,7 @@ class TestMergeSection:
|
|||
|
||||
def test_replace_existing_markers(self, tmp_path: Path):
|
||||
existing = tmp_path / "marked.md"
|
||||
existing.write_text(
|
||||
f"# Header\n\n{MARKER_START}\nold content\n{MARKER_END}\n\n# Footer"
|
||||
)
|
||||
existing.write_text(f"# Header\n\n{MARKER_START}\nold content\n{MARKER_END}\n\n# Footer")
|
||||
result = _merge_section(existing, f"{MARKER_START}\nnew content\n{MARKER_END}")
|
||||
assert "new content" in result
|
||||
assert "old content" not in result
|
||||
|
|
@ -204,7 +200,7 @@ class TestCursorWriter:
|
|||
def test_creates_mdc_with_frontmatter(self, tmp_path: Path):
|
||||
writer = CursorMemoryWriter(project_path=tmp_path)
|
||||
entries = _make_entries(3)
|
||||
result = writer.export(entries, dry_run=False)
|
||||
writer.export(entries, dry_run=False)
|
||||
|
||||
mdc_path = tmp_path / ".cursor" / "rules" / "headroom-memory.mdc"
|
||||
assert mdc_path.exists()
|
||||
|
|
@ -252,7 +248,7 @@ class TestCodexWriter:
|
|||
def test_export(self, tmp_path: Path):
|
||||
writer = CodexMemoryWriter(project_path=tmp_path)
|
||||
entries = _make_entries(3)
|
||||
result = writer.export(entries, dry_run=False)
|
||||
writer.export(entries, dry_run=False)
|
||||
|
||||
agents_md = tmp_path / "AGENTS.md"
|
||||
assert agents_md.exists()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue