learningcircuit--local-deep-research/tests/benchmarks/test_datasets_utils.py

"""
Tests for benchmarks/datasets/utils.py

Tests cover:
- derive_key function for key derivation
- decrypt function for XOR decryption
- get_known_answer_map function
"""

import base64


class TestDeriveKey:
    """Tests for the derive_key function."""

    def test_derives_consistent_key(self):
        """Test that same password produces same key."""
        from local_deep_research.benchmarks.datasets.utils import derive_key

        key1 = derive_key("password123", 32)
        key2 = derive_key("password123", 32)

        assert key1 == key2

    def test_different_passwords_produce_different_keys(self):
        """Test that different passwords produce different keys."""
        from local_deep_research.benchmarks.datasets.utils import derive_key

        key1 = derive_key("password1", 32)
        key2 = derive_key("password2", 32)

        assert key1 != key2

    def test_key_has_requested_length(self):
        """Test that key has the requested length."""
        from local_deep_research.benchmarks.datasets.utils import derive_key

        for length in [16, 32, 64, 100]:
            key = derive_key("password", length)
            assert len(key) == length

    def test_handles_short_length(self):
        """Test handling of short key lengths."""
        from local_deep_research.benchmarks.datasets.utils import derive_key

        key = derive_key("password", 8)
        assert len(key) == 8

    def test_handles_long_length(self):
        """Test handling of long key lengths (requires key repetition)."""
        from local_deep_research.benchmarks.datasets.utils import derive_key

        key = derive_key("password", 128)
        assert len(key) == 128

    def test_empty_password(self):
        """Test handling of empty password."""
        from local_deep_research.benchmarks.datasets.utils import derive_key

        key = derive_key("", 32)
        assert len(key) == 32


class TestDecrypt:
    """Tests for the decrypt function."""

    def test_returns_short_strings_unchanged(self):
        """Test that strings shorter than 8 chars are returned unchanged."""
        from local_deep_research.benchmarks.datasets.utils import decrypt

        result = decrypt("short", "password")
        assert result == "short"

    def test_returns_non_base64_unchanged(self):
        """Test that non-base64 strings are returned unchanged."""
        from local_deep_research.benchmarks.datasets.utils import decrypt

        # Contains characters not in base64 alphabet
        result = decrypt("not@base64!string", "password")
        assert result == "not@base64!string"

    def test_returns_non_string_unchanged(self):
        """Test that non-string input is returned unchanged."""
        from local_deep_research.benchmarks.datasets.utils import decrypt

        result = decrypt(12345, "password")
        assert result == 12345

    def test_known_encryption_decryption(self):
        """Test decryption with known encrypted value."""
        from local_deep_research.benchmarks.datasets.utils import (
            decrypt,
            derive_key,
        )

        # Create a known encrypted value
        plaintext = "Hello World test message"
        password = "test_password"

        # Encrypt
        key = derive_key(password, len(plaintext.encode()))
        encrypted = bytes(
            a ^ b for a, b in zip(plaintext.encode(), key, strict=False)
        )
        ciphertext_b64 = base64.b64encode(encrypted).decode()

        # Decrypt should return original
        result = decrypt(ciphertext_b64, password)
        assert result == plaintext

    def test_invalid_base64_returns_original(self):
        """Test that invalid base64 returns original string."""
        from local_deep_research.benchmarks.datasets.utils import decrypt

        # Valid base64 chars but invalid padding/encoding
        invalid = "AAAAAAAA"  # Valid chars but may not decode properly
        result = decrypt(invalid, "password")
        # Should return something (either original or decrypted)
        assert isinstance(result, str)

    def test_decryption_with_hardcoded_key_fallback(self):
        """Test that hardcoded key fallback is attempted."""
        from local_deep_research.benchmarks.datasets.utils import decrypt

        # Create encrypted with hardcoded BrowseComp key
        plaintext = "This is a test message with spaces"
        password = "MHGGF2022!"

        from local_deep_research.benchmarks.datasets.utils import derive_key

        key = derive_key(password, len(plaintext.encode()))
        encrypted = bytes(
            a ^ b for a, b in zip(plaintext.encode(), key, strict=False)
        )
        ciphertext_b64 = base64.b64encode(encrypted).decode()

        # Should decrypt using hardcoded key as fallback
        result = decrypt(ciphertext_b64, "wrong_password")
        assert result == plaintext


class TestGetKnownAnswerMap:
    """Tests for the get_known_answer_map function."""

    def test_returns_dictionary(self):
        """Test that function returns a dictionary."""
        from local_deep_research.benchmarks.datasets.utils import (
            get_known_answer_map,
        )

        result = get_known_answer_map()
        assert isinstance(result, dict)

    def test_contains_known_mappings(self):
        """Test that known mappings are present."""
        from local_deep_research.benchmarks.datasets.utils import (
            get_known_answer_map,
        )

        result = get_known_answer_map()

        # Check known mappings exist
        assert "dFoTn+K+bcdyWg==" in result
        assert result["dFoTn+K+bcdyWg=="] == "Tooth Rock"

        assert "ERFIwA==" in result
        assert result["ERFIwA=="] == "1945"

    def test_values_are_strings(self):
        """Test that all values are strings."""
        from local_deep_research.benchmarks.datasets.utils import (
            get_known_answer_map,
        )

        result = get_known_answer_map()

        for key, value in result.items():
            assert isinstance(key, str)
            assert isinstance(value, str)


class TestDatasetRegistry:
    """Tests for the DatasetRegistry class."""

    def test_registry_has_register_method(self):
        """Test that registry has register method."""
        from local_deep_research.benchmarks.datasets.base import DatasetRegistry

        assert hasattr(DatasetRegistry, "register")
        assert callable(DatasetRegistry.register)

    def test_registry_has_get_dataset_class_method(self):
        """Test that registry has get_dataset_class method."""
        from local_deep_research.benchmarks.datasets.base import DatasetRegistry

        assert hasattr(DatasetRegistry, "get_dataset_class")
        assert callable(DatasetRegistry.get_dataset_class)

    def test_registry_has_create_dataset_method(self):
        """Test that registry has create_dataset method."""
        from local_deep_research.benchmarks.datasets.base import DatasetRegistry

        assert hasattr(DatasetRegistry, "create_dataset")
        assert callable(DatasetRegistry.create_dataset)

    def test_registry_has_get_available_datasets(self):
        """Test that registry has get_available_datasets method."""
        from local_deep_research.benchmarks.datasets.base import DatasetRegistry

        assert hasattr(DatasetRegistry, "get_available_datasets")
        result = DatasetRegistry.get_available_datasets()
        assert isinstance(result, list)

    def test_registry_has_load_dataset_method(self):
        """Test that registry has load_dataset method."""
        from local_deep_research.benchmarks.datasets.base import DatasetRegistry

        assert hasattr(DatasetRegistry, "load_dataset")
        assert callable(DatasetRegistry.load_dataset)

    def test_get_unknown_dataset_raises(self):
        """Test that getting unknown dataset raises ValueError."""
        from local_deep_research.benchmarks.datasets.base import DatasetRegistry
        import pytest

        with pytest.raises(ValueError, match="Unknown dataset"):
            DatasetRegistry.get_dataset_class("nonexistent_dataset_xyz")


class TestBenchmarkDatasetBase:
    """Tests for the BenchmarkDataset abstract base class."""

    def test_benchmark_dataset_is_abstract(self):
        """Test that BenchmarkDataset is abstract."""
        from local_deep_research.benchmarks.datasets.base import (
            BenchmarkDataset,
        )
        from abc import ABC

        assert issubclass(BenchmarkDataset, ABC)

    def test_benchmark_dataset_has_abstract_methods(self):
        """Test that BenchmarkDataset has required abstract methods."""
        from local_deep_research.benchmarks.datasets.base import (
            BenchmarkDataset,
        )

        # Check for abstract methods
        assert hasattr(BenchmarkDataset, "get_dataset_info")
        assert hasattr(BenchmarkDataset, "get_default_dataset_path")
        assert hasattr(BenchmarkDataset, "process_example")

    def test_benchmark_dataset_has_utility_methods(self):
        """Test that BenchmarkDataset has utility methods."""
        from local_deep_research.benchmarks.datasets.base import (
            BenchmarkDataset,
        )

        assert hasattr(BenchmarkDataset, "load")
        assert hasattr(BenchmarkDataset, "get_examples")
        assert hasattr(BenchmarkDataset, "get_example")
        assert hasattr(BenchmarkDataset, "get_question")
        assert hasattr(BenchmarkDataset, "get_answer")