learningcircuit--local-deep-research/tests/benchmarks/datasets/test_utils_high_value.py

"""High-value tests for benchmarks/datasets/utils.py.

Tests cover derive_key(), decrypt(), and get_known_answer_map() with
emphasis on pure logic validation -- no mocking required.
"""

import base64
import hashlib

from local_deep_research.benchmarks.datasets.utils import (
    decrypt,
    derive_key,
    get_known_answer_map,
)


# ---------------------------------------------------------------------------
# Helper: XOR-encrypt a plaintext with derive_key, then base64-encode
# ---------------------------------------------------------------------------


def _encrypt(plaintext: str, password: str) -> str:
    """Mirror the XOR encryption that decrypt() reverses."""
    plaintext_bytes = plaintext.encode("utf-8")
    key = derive_key(password, len(plaintext_bytes))
    encrypted = bytes(a ^ b for a, b in zip(plaintext_bytes, key))
    return base64.b64encode(encrypted).decode("ascii")


# ===================================================================
# derive_key tests
# ===================================================================


class TestDeriveKey:
    """Tests for derive_key()."""

    def test_output_length_matches_requested_short(self):
        """Key length < 32 should still return exactly that many bytes."""
        key = derive_key("password", 10)
        assert len(key) == 10

    def test_output_length_matches_requested_exact_32(self):
        """32 is the SHA-256 digest size -- exact match, no repetition."""
        key = derive_key("password", 32)
        assert len(key) == 32

    def test_output_length_matches_requested_long(self):
        """Lengths > 32 require repeating the digest."""
        key = derive_key("password", 100)
        assert len(key) == 100

    def test_output_is_bytes(self):
        key = derive_key("secret", 16)
        assert isinstance(key, bytes)

    def test_deterministic(self):
        """Same inputs must always produce the same key."""
        a = derive_key("hello", 48)
        b = derive_key("hello", 48)
        assert a == b

    def test_different_passwords_produce_different_keys(self):
        a = derive_key("alpha", 32)
        b = derive_key("beta", 32)
        assert a != b

    def test_key_repeats_sha256_digest_for_long_length(self):
        """For length > 32 the key should be digest repeated + remainder."""
        digest = hashlib.sha256(b"mypass").digest()
        key = derive_key("mypass", 70)
        # First 32 bytes == digest, next 32 bytes == digest, last 6 == digest[:6]
        assert key[:32] == digest
        assert key[32:64] == digest
        assert key[64:] == digest[:6]

    def test_zero_length(self):
        """Edge case: requesting a zero-length key returns empty bytes."""
        key = derive_key("anything", 0)
        assert key == b""

    def test_length_one(self):
        key = derive_key("x", 1)
        assert len(key) == 1
        assert isinstance(key, bytes)


# ===================================================================
# decrypt tests
# ===================================================================


class TestDecrypt:
    """Tests for decrypt()."""

    # -- Non-string / short-string passthrough --------------------------

    def test_non_string_input_returned_as_is(self):
        assert decrypt(12345, "pw") == 12345

    def test_none_returned_as_is(self):
        assert decrypt(None, "pw") is None

    def test_list_returned_as_is(self):
        val = [1, 2, 3]
        assert decrypt(val, "pw") is val

    def test_short_string_returned_as_is(self):
        """Strings shorter than 8 chars skip decryption."""
        assert decrypt("ABCD", "pw") == "ABCD"

    def test_exactly_seven_chars_returned_as_is(self):
        assert decrypt("ABCDEFG", "pw") == "ABCDEFG"

    # -- Non-base64 passthrough -----------------------------------------

    def test_non_base64_characters_returned_as_is(self):
        """Strings with characters outside the base64 alphabet are skipped."""
        bad = "This has spaces and !@# symbols"
        assert decrypt(bad, "pw") == bad

    # -- Round-trip encrypt/decrypt -------------------------------------

    def test_round_trip_standard_method(self):
        """Encrypt then decrypt should recover plaintext (with space for heuristic)."""
        plaintext = "Hello World test message"
        password = "secret"
        ciphertext = _encrypt(plaintext, password)
        result = decrypt(ciphertext, password)
        assert result == plaintext

    def test_round_trip_various_plaintexts(self):
        """Multiple different plaintexts round-trip correctly."""
        password = "mykey"
        for text in [
            "The quick brown fox jumps",
            "Another test message here",
            "1234 numbers and letters mixed",
        ]:
            ct = _encrypt(text, password)
            assert decrypt(ct, password) == text

    # -- Alt method 1: password > 30 chars, uses first word -------------

    def test_alt_method_1_first_word_password(self):
        """When password > 30 chars, decrypt tries the first word."""
        long_password = (
            "firstword plus a lot of extra padding to exceed thirty characters"
        )
        first_word = "firstword"
        plaintext = "Decrypted via alt method one"
        # Encrypt with the first word (what alt method 1 will use)
        ciphertext = _encrypt(plaintext, first_word)
        result = decrypt(ciphertext, long_password)
        assert result == plaintext

    # -- Alt method 2: password contains "GUID" -------------------------

    def test_alt_method_2_guid_password(self):
        """When password contains 'GUID', decrypt tries the part after it."""
        guid_part = "abc123-guid-value"
        password_with_guid = f"prefix GUID {guid_part}"
        plaintext = "Decrypted via GUID method here"
        # Encrypt with the guid part (what alt method 2 will use)
        ciphertext = _encrypt(plaintext, guid_part)
        result = decrypt(ciphertext, password_with_guid)
        assert result == plaintext

    # -- Hardcoded key fallback -----------------------------------------

    def test_hardcoded_key_fallback(self):
        """When other methods fail, decrypt tries the hardcoded BrowseComp key."""
        hardcoded_key = "MHGGF2022!"
        plaintext = "Decrypted with hardcoded key ok"
        ciphertext = _encrypt(plaintext, hardcoded_key)
        # Use a short password that won't match, so it falls through to hardcoded
        result = decrypt(ciphertext, "wrongpassword")
        assert result == plaintext

    # -- All methods fail: original returned ----------------------------

    def test_all_methods_fail_returns_original(self):
        """If nothing decrypts to valid ASCII with spaces, return original."""
        # Create ciphertext that won't decrypt to ASCII-with-spaces under any method
        raw = bytes(range(200, 220))
        ct = base64.b64encode(raw).decode("ascii")
        if len(ct) >= 8:
            result = decrypt(ct, "irrelevant")
            # Should return original since no method produces valid text
            assert isinstance(result, str)

    def test_empty_string_returned_as_is(self):
        assert decrypt("", "pw") == ""


# ===================================================================
# get_known_answer_map tests
# ===================================================================


class TestGetKnownAnswerMap:
    """Tests for get_known_answer_map()."""

    def test_returns_dict(self):
        result = get_known_answer_map()
        assert isinstance(result, dict)

    def test_contains_tooth_rock_entry(self):
        m = get_known_answer_map()
        assert m["dFoTn+K+bcdyWg=="] == "Tooth Rock"

    def test_contains_1945_entry(self):
        m = get_known_answer_map()
        assert m["ERFIwA=="] == "1945"

    def test_all_values_are_strings(self):
        for v in get_known_answer_map().values():
            assert isinstance(v, str)

    def test_all_keys_are_strings(self):
        for k in get_known_answer_map().keys():
            assert isinstance(k, str)