chopratejas--headroom/tests/test_compression/test_masks.py

"""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