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

"""
Tests targeting specific uncovered lines in xbench_deepsearch.py.

Covers:
- Decrypt failure in load_data (exception in base64 decode/XOR)
- Encrypted prompt in _load_from_url
- Decrypt failure in _load_from_url
- _load_from_url returns every row (sampling happens in load())
"""

import base64
from unittest.mock import MagicMock, Mock, patch

MODULE = "local_deep_research.benchmarks.datasets.xbench_deepsearch"


def _make_dataset():
    from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
        XBenchDeepSearchDataset,
    )

    return XBenchDeepSearchDataset()


# ---------------------------------------------------------------------------
# Lines 133-138: Decrypt failure in load_data
# ---------------------------------------------------------------------------
class TestLoadDataDecryptFailure:
    def test_decrypt_failure_logs_warning_and_keeps_original(self):
        """When XOR decrypt produces invalid UTF-8, exception is caught
        and original (encrypted) value is used."""
        ds = _make_dataset()

        # Create a prompt that looks like base64 but when decrypted gives invalid UTF-8
        # Use bytes that are valid base64 but XOR to non-UTF-8 sequences
        bad_bytes = bytes([0xFF, 0xFE, 0xFD, 0xFC])  # Invalid UTF-8
        key = "k"
        key_bytes = key.encode("utf-8")
        # XOR to create "encrypted" form
        encrypted = bytes(
            [
                bad_bytes[i] ^ key_bytes[i % len(key_bytes)]
                for i in range(len(bad_bytes))
            ]
        )
        encoded_prompt = base64.b64encode(encrypted).decode("utf-8")

        mock_item = {
            "id": "fail_1",
            "prompt": encoded_prompt,
            "answer": "plain answer",
            "canary": key,
            "reference_steps": "",
        }
        mock_dataset = [mock_item]
        mock_load_dataset = Mock(return_value=mock_dataset)

        with patch.dict(
            "sys.modules", {"datasets": Mock(load_dataset=mock_load_dataset)}
        ):
            with patch(f"{MODULE}.logger") as mock_logger:
                result = ds.load_data()

        assert len(result) == 1
        # The warning should have been called
        mock_logger.warning.assert_called()


# ---------------------------------------------------------------------------
# Encrypted prompt in _load_from_url
# ---------------------------------------------------------------------------
class TestLoadFromUrlEncrypted:
    def test_url_decrypts_base64_prompt(self):
        """_load_from_url decrypts base64-encoded prompts via XOR."""
        ds = _make_dataset()

        plaintext = b"Real question from URL"
        key = "urlkey"
        key_bytes = key.encode("utf-8")
        xored = bytes(
            [
                plaintext[i] ^ key_bytes[i % len(key_bytes)]
                for i in range(len(plaintext))
            ]
        )
        encoded_prompt = base64.b64encode(xored).decode("utf-8")

        mock_df = MagicMock()
        row = MagicMock()
        row.get = lambda k, default="": {
            "id": "url_enc_1",
            "prompt": encoded_prompt,
            "answer": "plain answer",
            "canary": key,
            "reference_steps": "",
        }.get(k, default)
        mock_df.iterrows.return_value = [(0, row)]

        mock_pd = Mock()
        mock_pd.read_parquet.return_value = mock_df
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger"):
                result = ds._load_from_url()

        assert len(result) == 1
        assert result[0]["problem"] == "Real question from URL"


# ---------------------------------------------------------------------------
# Decrypt failure in _load_from_url
# ---------------------------------------------------------------------------
class TestLoadFromUrlDecryptFailure:
    def test_url_decrypt_failure_keeps_original(self):
        """When decrypt fails in _load_from_url, original value is used."""
        ds = _make_dataset()

        bad_bytes = bytes([0xFF, 0xFE, 0xFD])
        key = "x"
        key_bytes = key.encode("utf-8")
        encrypted = bytes(
            [
                bad_bytes[i] ^ key_bytes[i % len(key_bytes)]
                for i in range(len(bad_bytes))
            ]
        )
        encoded_prompt = base64.b64encode(encrypted).decode("utf-8")

        mock_df = MagicMock()
        row = MagicMock()
        row.get = lambda k, default="": {
            "id": "url_fail_1",
            "prompt": encoded_prompt,
            "answer": "plain",
            "canary": key,
            "reference_steps": "",
        }.get(k, default)
        mock_df.iterrows.return_value = [(0, row)]

        mock_pd = Mock()
        mock_pd.read_parquet.return_value = mock_df
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger") as mock_logger:
                result = ds._load_from_url()

        assert len(result) == 1
        mock_logger.warning.assert_called()


# ---------------------------------------------------------------------------
# _load_from_url returns every row (sampling happens in load())
# ---------------------------------------------------------------------------
class TestLoadFromUrlLoadsAll:
    def test_url_loads_all_rows(self):
        """_load_from_url returns the full dataset without sampling."""
        ds = _make_dataset()

        mock_df = MagicMock()
        rows = []
        for i in range(10):
            row = MagicMock()
            row.get = lambda k, default="", idx=i: {
                "id": f"url_q{idx}",
                "prompt": f"Question {idx}",
                "answer": f"Answer {idx}",
                "canary": "",
                "reference_steps": "",
            }.get(k, default)
            rows.append((i, row))
        mock_df.iterrows.return_value = rows

        mock_pd = Mock()
        mock_pd.read_parquet.return_value = mock_df
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger"):
                result = ds._load_from_url()

        assert len(result) == 10
        assert [q["id"] for q in result] == [f"url_q{i}" for i in range(10)]
        assert result[0]["problem"] == "Question 0"