mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
31aa72c885
- Add headroom.compression module with UniversalCompressor - ML-based content detection using Magika (JSON, code, logs, text) - Structure-preserving compression via handler protocol - JSON handler: preserves keys, brackets, high-entropy values (UUIDs) - Code handler: preserves imports, signatures, types (tree-sitter AST) - Entropy-based preservation for identifiers and hashes - CCR integration for reversible compression - Comprehensive test suite with LLM eval tests - Add docs/compression.md with full API documentation
272 lines
8.6 KiB
Python
272 lines
8.6 KiB
Python
"""Tests for structure mask system."""
|
|
|
|
import pytest
|
|
|
|
from headroom.compression.masks import (
|
|
EntropyScore,
|
|
MaskSpan,
|
|
StructureMask,
|
|
apply_mask_to_text,
|
|
compute_entropy_mask,
|
|
mask_to_spans,
|
|
)
|
|
|
|
|
|
class TestStructureMask:
|
|
"""Tests for StructureMask class."""
|
|
|
|
def test_create_mask(self):
|
|
"""Test basic mask creation."""
|
|
tokens = ["a", "b", "c", "d"]
|
|
mask = [True, False, False, True]
|
|
|
|
sm = StructureMask(tokens=tokens, mask=mask)
|
|
|
|
assert len(sm.tokens) == 4
|
|
assert len(sm.mask) == 4
|
|
assert sm.structural_count == 2
|
|
assert sm.compressible_count == 2
|
|
|
|
def test_mask_length_mismatch_raises(self):
|
|
"""Test that mismatched lengths raise ValueError."""
|
|
tokens = ["a", "b", "c"]
|
|
mask = [True, False] # Wrong length
|
|
|
|
with pytest.raises(ValueError, match="must match"):
|
|
StructureMask(tokens=tokens, mask=mask)
|
|
|
|
def test_preservation_ratio(self):
|
|
"""Test preservation ratio calculation."""
|
|
tokens = list("abcdefghij") # 10 tokens
|
|
mask = [True, True, False, False, False, False, False, False, False, False]
|
|
|
|
sm = StructureMask(tokens=tokens, mask=mask)
|
|
|
|
assert sm.preservation_ratio == 0.2 # 2/10
|
|
|
|
def test_empty_mask(self):
|
|
"""Test creating empty mask (all compressible)."""
|
|
tokens = list("hello")
|
|
sm = StructureMask.empty(tokens)
|
|
|
|
assert all(not m for m in sm.mask)
|
|
assert sm.preservation_ratio == 0.0
|
|
|
|
def test_full_mask(self):
|
|
"""Test creating full mask (all preserved)."""
|
|
tokens = list("hello")
|
|
sm = StructureMask.full(tokens)
|
|
|
|
assert all(m for m in sm.mask)
|
|
assert sm.preservation_ratio == 1.0
|
|
|
|
def test_get_structural_tokens(self):
|
|
"""Test extracting structural tokens."""
|
|
tokens = ["def", " ", "foo", "(", ")", ":"]
|
|
mask = [True, False, True, True, True, True]
|
|
|
|
sm = StructureMask(tokens=tokens, mask=mask)
|
|
structural = sm.get_structural_tokens()
|
|
|
|
assert structural == ["def", "foo", "(", ")", ":"]
|
|
|
|
def test_get_compressible_tokens(self):
|
|
"""Test extracting compressible tokens."""
|
|
tokens = ["def", " ", "foo", "(", ")", ":"]
|
|
mask = [True, False, True, True, True, True]
|
|
|
|
sm = StructureMask(tokens=tokens, mask=mask)
|
|
compressible = sm.get_compressible_tokens()
|
|
|
|
assert compressible == [" "]
|
|
|
|
def test_union_masks(self):
|
|
"""Test union of two masks."""
|
|
tokens = list("abcd")
|
|
mask1 = StructureMask(tokens=tokens, mask=[True, False, False, False])
|
|
mask2 = StructureMask(tokens=tokens, mask=[False, False, True, False])
|
|
|
|
result = mask1.union(mask2)
|
|
|
|
assert result.mask == [True, False, True, False]
|
|
|
|
def test_union_different_lengths_raises(self):
|
|
"""Test that union of different length masks raises."""
|
|
mask1 = StructureMask(tokens=["a", "b"], mask=[True, False])
|
|
mask2 = StructureMask(tokens=["a", "b", "c"], mask=[True, False, True])
|
|
|
|
with pytest.raises(ValueError, match="different lengths"):
|
|
mask1.union(mask2)
|
|
|
|
def test_intersection_masks(self):
|
|
"""Test intersection of two masks."""
|
|
tokens = list("abcd")
|
|
mask1 = StructureMask(tokens=tokens, mask=[True, True, False, False])
|
|
mask2 = StructureMask(tokens=tokens, mask=[True, False, True, False])
|
|
|
|
result = mask1.intersection(mask2)
|
|
|
|
assert result.mask == [True, False, False, False]
|
|
|
|
|
|
class TestMaskToSpans:
|
|
"""Tests for mask_to_spans function."""
|
|
|
|
def test_simple_spans(self):
|
|
"""Test converting mask to spans."""
|
|
tokens = list("abcdef")
|
|
mask = StructureMask(
|
|
tokens=tokens,
|
|
mask=[True, True, True, False, False, False],
|
|
)
|
|
|
|
spans = mask_to_spans(mask)
|
|
|
|
assert len(spans) == 2
|
|
assert spans[0] == MaskSpan(start=0, end=3, is_structural=True)
|
|
assert spans[1] == MaskSpan(start=3, end=6, is_structural=False)
|
|
|
|
def test_alternating_spans(self):
|
|
"""Test mask with alternating regions."""
|
|
tokens = list("abcdef")
|
|
mask = StructureMask(
|
|
tokens=tokens,
|
|
mask=[True, False, True, False, True, False],
|
|
)
|
|
|
|
spans = mask_to_spans(mask)
|
|
|
|
assert len(spans) == 6 # Each token is its own span
|
|
|
|
def test_empty_mask(self):
|
|
"""Test empty mask produces no spans."""
|
|
mask = StructureMask(tokens=[], mask=[])
|
|
spans = mask_to_spans(mask)
|
|
|
|
assert spans == []
|
|
|
|
def test_span_length(self):
|
|
"""Test span length property."""
|
|
span = MaskSpan(start=5, end=15, is_structural=True)
|
|
assert span.length == 10
|
|
|
|
|
|
class TestEntropyScore:
|
|
"""Tests for entropy-based preservation."""
|
|
|
|
def test_high_entropy_uuid(self):
|
|
"""Test that UUIDs have high entropy."""
|
|
uuid = "8f14e45f-ceea-4123-8f14-e45fceea4123"
|
|
score = EntropyScore.compute(uuid, threshold=0.8)
|
|
|
|
assert score.value > 0.8
|
|
assert score.should_preserve is True
|
|
|
|
def test_low_entropy_repeated(self):
|
|
"""Test that repeated text has low entropy."""
|
|
text = "aaaaaaaaaaaaaaaa"
|
|
score = EntropyScore.compute(text, threshold=0.5)
|
|
|
|
assert score.value < 0.3
|
|
assert score.should_preserve is False
|
|
|
|
def test_normal_text_entropy(self):
|
|
"""Test normal text entropy."""
|
|
text = "The quick brown fox"
|
|
score = EntropyScore.compute(text, threshold=0.85)
|
|
|
|
# Normal diverse text has high entropy (no repetition)
|
|
assert 0.5 < score.value <= 1.0
|
|
|
|
def test_empty_text(self):
|
|
"""Test empty text."""
|
|
score = EntropyScore.compute("", threshold=0.5)
|
|
|
|
assert score.value == 0.0
|
|
assert score.should_preserve is False
|
|
|
|
def test_custom_threshold(self):
|
|
"""Test custom threshold."""
|
|
text = "abc123xyz" # Moderate entropy
|
|
|
|
high_threshold = EntropyScore.compute(text, threshold=0.95)
|
|
low_threshold = EntropyScore.compute(text, threshold=0.5)
|
|
|
|
# Same value, different preservation decisions
|
|
assert high_threshold.value == low_threshold.value
|
|
assert (
|
|
high_threshold.should_preserve != low_threshold.should_preserve
|
|
or high_threshold.value >= 0.95
|
|
or high_threshold.value < 0.5
|
|
)
|
|
|
|
|
|
class TestComputeEntropyMask:
|
|
"""Tests for compute_entropy_mask function."""
|
|
|
|
def test_preserves_uuids(self):
|
|
"""Test that UUIDs are preserved."""
|
|
tokens = ["user", ":", " ", "8f14e45f-ceea-4123-8f14-e45fceea4123"]
|
|
mask = compute_entropy_mask(tokens, threshold=0.8)
|
|
|
|
# Only the UUID token should be preserved
|
|
assert mask.mask[0] is False # "user"
|
|
assert mask.mask[1] is False # ":"
|
|
assert mask.mask[2] is False # " "
|
|
assert mask.mask[3] is True # UUID
|
|
|
|
def test_short_tokens_not_checked(self):
|
|
"""Test that short tokens are not checked for entropy."""
|
|
tokens = ["ab", "cd", "ef"]
|
|
mask = compute_entropy_mask(tokens, min_token_length=10)
|
|
|
|
# All tokens too short to check
|
|
assert all(not m for m in mask.mask)
|
|
|
|
def test_metadata_contains_threshold(self):
|
|
"""Test that metadata contains threshold."""
|
|
tokens = ["test"]
|
|
mask = compute_entropy_mask(tokens, threshold=0.9)
|
|
|
|
assert mask.metadata["source"] == "entropy"
|
|
assert mask.metadata["threshold"] == 0.9
|
|
|
|
|
|
class TestApplyMaskToText:
|
|
"""Tests for apply_mask_to_text function."""
|
|
|
|
def test_preserves_structural(self):
|
|
"""Test that structural regions are preserved."""
|
|
text = "def foo(): pass"
|
|
tokens = list(text)
|
|
mask = StructureMask(
|
|
tokens=tokens,
|
|
# Preserve "def foo():" (first 10 chars)
|
|
mask=[True] * 10 + [False] * 5,
|
|
)
|
|
|
|
def mock_compress(s: str) -> str:
|
|
return "[C]"
|
|
|
|
result = apply_mask_to_text(text, mask, mock_compress)
|
|
|
|
assert result.startswith("def foo():")
|
|
assert "[C]" in result
|
|
|
|
def test_compresses_non_structural(self):
|
|
"""Test that non-structural regions are compressed."""
|
|
text = "aaa bbb ccc"
|
|
tokens = list(text)
|
|
mask = StructureMask(
|
|
tokens=tokens,
|
|
mask=[True, True, True, False, False, False, False, True, True, True, True],
|
|
)
|
|
|
|
def mock_compress(s: str) -> str:
|
|
return "X"
|
|
|
|
result = apply_mask_to_text(text, mask, mock_compress)
|
|
|
|
# "aaa" preserved, " bbb " compressed to "X", "ccc" preserved
|
|
assert "aaa" in result
|
|
assert "ccc" in result
|