"""Tests for structure mask system.""" import pytest from headroom.compression.masks import ( EntropyScore, MaskSpan, StructureMask, apply_mask_to_text, compute_entropy_mask, compute_entropy_mask_for_content, 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 TestComputeEntropyMaskForContent: """Tests for compute_entropy_mask_for_content (SEC-01 regression). The character-level path (`compute_entropy_mask(list(content))`) is a silent no-op on plain text because every single-character token is below min_token_length. The content-level helper must restore preservation by scoring whole words and mapping them back to character positions. """ def test_char_level_tokenization_is_inert(self): """Regression: char tokens never reach min length -> nothing preserved.""" secret = "Zx9Kq3Wm7Pv2Lr8Nt4Bc6Df1Gh5Jy" # gitleaks:allow synthetic test fixture char_mask = compute_entropy_mask(list(f"k={secret}"), threshold=0.85) # This is the bug the fix routes around: zero preservation. assert sum(char_mask.mask) == 0 def test_high_entropy_word_char_range_preserved(self): """The full character span of a high-entropy word is marked True.""" secret = "Zx9Kq3Wm7Pv2Lr8Nt4Bc6Df1Gh5Jy" # gitleaks:allow synthetic test fixture content = f"prefix {secret} suffix" mask = compute_entropy_mask_for_content(content, threshold=0.85) start = content.index(secret) end = start + len(secret) assert all(mask.mask[start:end]) # secret preserved assert not any(mask.mask[:start]) # ordinary words not preserved assert not any(mask.mask[end:]) # trailing words not preserved assert len(mask.mask) == len(content) # char-aligned def test_short_words_not_preserved(self): """Short words are not scored regardless of entropy.""" mask = compute_entropy_mask_for_content("a b cd ef", threshold=0.5) assert sum(mask.mask) == 0 def test_metadata_marks_word_granularity(self): mask = compute_entropy_mask_for_content("plain words only", threshold=0.9) assert mask.metadata["source"] == "entropy" assert mask.metadata["threshold"] == 0.9 assert mask.metadata["granularity"] == "word" 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