/tests/benchmarks/datasets/test_xbench_deepsearch_coverage.py

"""
Coverage tests for benchmarks/datasets/xbench_deepsearch.py.

Targets the 74 missing lines:
- get_dataset_info / get_default_dataset_path
- load() with/without num_examples
- load_data() via datasets library – encrypted and plain prompts
- load_data() datasets.ImportError fallback to _load_from_url
- load_data() exception fallback to _load_from_url
- _load_from_url() success, exception (returns [])
- process_example()
- xor_decrypt
"""

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

MODULE = "local_deep_research.benchmarks.datasets.xbench_deepsearch"


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

    return XBenchDeepSearchDataset(**kwargs)


# ---------------------------------------------------------------------------
# Static helpers
# ---------------------------------------------------------------------------


class TestXorDecrypt:
    # test_xor_decrypt_roundtrip is defined first so it runs first and warms up
    # the expensive module import within the 60-second pytest-timeout window.
    def test_xor_decrypt_roundtrip(self):
        from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
            XBenchDeepSearchDataset,
        )

        plaintext = b"Hello, World!"
        key = "secret"
        encrypted = XBenchDeepSearchDataset.xor_decrypt(plaintext, key)
        decrypted = XBenchDeepSearchDataset.xor_decrypt(encrypted, key)
        assert decrypted == plaintext

    def test_xor_decrypt_basic(self):
        """Pure XOR decrypt: a key XOR'd twice returns the original data."""
        from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
            XBenchDeepSearchDataset,
        )

        plaintext = b"Hello, World!"
        key = "secret"
        encrypted = XBenchDeepSearchDataset.xor_decrypt(plaintext, key)
        # XOR is its own inverse
        decrypted = XBenchDeepSearchDataset.xor_decrypt(encrypted, key)
        assert decrypted == plaintext

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

        result = XBenchDeepSearchDataset.xor_decrypt(b"", "key")
        assert result == b""

    def test_xor_decrypt_key_wraps_around(self):
        """Key shorter than data: key bytes repeat (modulo)."""
        from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
            XBenchDeepSearchDataset,
        )

        data = b"ABCDEF"
        key = "AB"  # 2 bytes, data is 6 bytes – wraps 3 times
        result = XBenchDeepSearchDataset.xor_decrypt(data, key)
        assert len(result) == 6
        key_bytes = key.encode("utf-8")
        expected = bytes(data[i] ^ key_bytes[i % 2] for i in range(6))
        assert result == expected


# ---------------------------------------------------------------------------
# Dataset metadata
# ---------------------------------------------------------------------------


class TestDatasetInfo:
    def test_get_dataset_info_returns_required_keys(self):
        from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
            XBenchDeepSearchDataset,
        )

        info = XBenchDeepSearchDataset.get_dataset_info()
        assert info["id"] == "xbench_deepsearch"
        assert "name" in info
        assert "description" in info
        assert "url" in info

    def test_get_dataset_info(self):
        """get_dataset_info returns a dict with all expected keys and correct id."""
        from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
            XBenchDeepSearchDataset,
        )

        info = XBenchDeepSearchDataset.get_dataset_info()
        assert isinstance(info, dict)
        for key in ("id", "name", "description", "url"):
            assert key in info, f"Missing key: {key}"
        assert info["id"] == "xbench_deepsearch"
        assert "xbench" in info["url"].lower() or "DeepSearch" in info["url"]

    def test_get_default_dataset_path(self):
        """get_default_dataset_path returns a non-empty string."""
        from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
            XBenchDeepSearchDataset,
        )

        path = XBenchDeepSearchDataset.get_default_dataset_path()
        assert isinstance(path, str)
        assert len(path) > 0
        assert "xbench" in path.lower() or "DeepSearch" in path

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

        path = XBenchDeepSearchDataset.get_default_dataset_path()
        assert "/" in path


# ---------------------------------------------------------------------------
# process_example
# ---------------------------------------------------------------------------


class TestProcessExample:
    def test_process_example_adds_metadata(self):
        ds = _make_dataset()
        example = {"id": "q1", "problem": "What is X?", "answer": "Y"}
        result = ds.process_example(example)
        assert result["requires_deep_search"] is True
        assert result["expected_iterations"] == 4
        assert "evaluation_criteria" in result

    def test_process_example(self):
        """process_example adds requires_deep_search, expected_iterations, evaluation_criteria."""
        ds = _make_dataset()
        example = {
            "id": "q42",
            "problem": "Who invented calculus?",
            "answer": "Newton and Leibniz",
            "canary": "somekey",
        }
        result = ds.process_example(example)
        assert result["requires_deep_search"] is True
        assert result["expected_iterations"] == 4
        criteria = result["evaluation_criteria"]
        assert "accuracy" in criteria
        assert "completeness" in criteria
        assert "reasoning" in criteria
        assert "sources" in criteria

    def test_process_example_preserves_original_fields(self):
        ds = _make_dataset()
        example = {
            "id": "q2",
            "problem": "question",
            "answer": "ans",
            "canary": "key",
        }
        result = ds.process_example(example)
        assert result["id"] == "q2"
        assert result["problem"] == "question"
        assert result["answer"] == "ans"

    def test_evaluation_criteria_weights_sum_to_one(self):
        ds = _make_dataset()
        example = {"id": "q3", "problem": "Q", "answer": "A"}
        result = ds.process_example(example)
        total = sum(result["evaluation_criteria"].values())
        assert abs(total - 1.0) < 1e-9

    def test_process_example_does_not_mutate_input(self):
        """process_example returns a copy, original dict unchanged."""
        ds = _make_dataset()
        example = {"id": "q5", "problem": "P", "answer": "A"}
        original_keys = set(example.keys())
        ds.process_example(example)
        assert set(example.keys()) == original_keys


# ---------------------------------------------------------------------------
# load() method
# ---------------------------------------------------------------------------


class TestLoad:
    def test_load_returns_list_of_processed_examples(self):
        ds = _make_dataset()
        raw = [
            {"id": "1", "problem": "P1", "answer": "A1", "canary": ""},
            {"id": "2", "problem": "P2", "answer": "A2", "canary": ""},
        ]
        with patch.object(ds, "load_data", return_value=raw):
            result = ds.load()
        assert len(result) == 2
        assert result[0]["requires_deep_search"] is True

    def test_load_with_sampling(self):
        """load() samples when num_examples < total length."""
        ds = _make_dataset(num_examples=3, seed=42)
        raw = [
            {
                "id": str(i),
                "problem": f"Q{i}",
                "answer": f"A{i}",
                "canary": "",
            }
            for i in range(10)
        ]
        with patch.object(ds, "load_data", return_value=raw):
            result = ds.load()
        assert len(result) == 3
        for item in result:
            assert item["requires_deep_search"] is True

    def test_load_samples_when_num_examples_set(self):
        ds = _make_dataset(num_examples=3, seed=42)
        raw = [
            {
                "id": str(i),
                "problem": f"Q{i}",
                "answer": f"A{i}",
                "canary": "",
            }
            for i in range(10)
        ]
        with patch.object(ds, "load_data", return_value=raw):
            result = ds.load()
        assert len(result) == 3

    def test_load_no_sampling_when_num_examples_none(self):
        ds = _make_dataset()
        raw = [
            {
                "id": str(i),
                "problem": f"Q{i}",
                "answer": f"A{i}",
                "canary": "",
            }
            for i in range(5)
        ]
        with patch.object(ds, "load_data", return_value=raw):
            result = ds.load()
        assert len(result) == 5

    def test_load_full_integration(self):
        """load() calls load_data then process_example for each item."""
        ds = _make_dataset()
        raw = [
            {
                "id": "a",
                "problem": "Integration Q",
                "answer": "Integration A",
                "canary": "",
            },
        ]
        with patch.object(ds, "load_data", return_value=raw) as mock_load:
            with patch.object(
                ds, "process_example", wraps=ds.process_example
            ) as mock_process:
                result = ds.load()
        mock_load.assert_called_once()
        mock_process.assert_called_once_with(raw[0])
        assert len(result) == 1
        assert result[0]["requires_deep_search"] is True

    def test_load_caches_results(self):
        """A second load() returns the cached examples without reloading."""
        ds = _make_dataset()
        raw = [{"id": "1", "problem": "P", "answer": "A", "canary": ""}]
        with patch.object(ds, "load_data", return_value=raw) as mock_load:
            first = ds.load()
            second = ds.load()
        mock_load.assert_called_once()
        assert first is second

    def test_registry_load_dataset_respects_num_examples(self):
        """Regression test for #4451: DatasetRegistry.load_dataset() calls
        load() with no arguments, so sampling must honor the constructor's
        num_examples instead of loading the full dataset."""
        from local_deep_research.benchmarks.datasets.base import (
            DatasetRegistry,
        )
        from local_deep_research.benchmarks.datasets.xbench_deepsearch import (
            XBenchDeepSearchDataset,
        )

        raw = [
            {
                "id": str(i),
                "problem": f"Q{i}",
                "answer": f"A{i}",
                "canary": "",
            }
            for i in range(100)
        ]
        with patch.object(
            XBenchDeepSearchDataset, "load_data", return_value=raw
        ):
            result = DatasetRegistry.load_dataset(
                "xbench_deepsearch", num_examples=10, seed=None
            )
        assert len(result) == 10


# ---------------------------------------------------------------------------
# load_data() – datasets library available, plain text fields
# ---------------------------------------------------------------------------


class TestLoadDataPlainText:
    def test_load_via_datasets_library(self):
        """load_data uses datasets.load_dataset when available and returns formatted items."""
        ds = _make_dataset()
        mock_item = {
            "id": "plain_1",
            "prompt": "What is the speed of light?",
            "answer": "299,792,458 m/s",
            "canary": "testkey",
            "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"):
                result = ds.load_data()
        assert len(result) == 1
        assert result[0]["problem"] == "What is the speed of light?"
        assert result[0]["answer"] == "299,792,458 m/s"

    def test_load_data_plain_text_prompt_and_answer(self):
        ds = _make_dataset()
        mock_item = {
            "id": "plain_1",
            "prompt": "What is the speed of light?",
            "answer": "299,792,458 m/s",
            "canary": "testkey",
            "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"):
                result = ds.load_data()
        assert len(result) == 1
        assert result[0]["problem"] == "What is the speed of light?"
        assert result[0]["answer"] == "299,792,458 m/s"

    def test_load_data_uses_default_path_when_none(self):
        ds = _make_dataset()
        mock_dataset = []
        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"):
                result = ds.load_data(dataset_path=None)
        mock_load_dataset.assert_called_once_with(
            "xbench/DeepSearch", split="train"
        )
        assert result == []

    def test_load_data_formatted_item_has_required_keys(self):
        ds = _make_dataset()
        mock_item = {
            "id": "test_id",
            "prompt": "Some plain question",
            "answer": "Some plain answer",
            "canary": "k",
            "reference_steps": "step1",
        }
        mock_load_dataset = Mock(return_value=[mock_item])
        with patch.dict(
            "sys.modules", {"datasets": Mock(load_dataset=mock_load_dataset)}
        ):
            with patch(f"{MODULE}.logger"):
                result = ds.load_data()
        item = result[0]
        for key in ("id", "problem", "answer", "reference_steps", "canary"):
            assert key in item, f"Missing key: {key}"
        assert item["id"] == "test_id"


# ---------------------------------------------------------------------------
# load_data() – encrypted fields (base64-like prompt)
# ---------------------------------------------------------------------------


class TestLoadDataEncryptedFields:
    def test_load_data_decrypts_base64_encoded_prompt(self):
        ds = _make_dataset()
        key = "mykey"
        plaintext = b"This is the real question"
        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_item = {
            "id": "enc_1",
            "prompt": encoded_prompt,
            "answer": "some 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"):
                result = ds.load_data()
        assert result[0]["problem"] == "This is the real question"


# ---------------------------------------------------------------------------
# load_data() – datasets ImportError falls back to _load_from_url
# ---------------------------------------------------------------------------


class TestLoadDatasetsImportErrorFallback:
    def test_load_datasets_import_error_fallback(self):
        """When 'datasets' library is not installed, _load_from_url is called."""
        ds = _make_dataset()
        fallback_data = [
            {"id": "url_1", "problem": "Q", "answer": "A", "canary": ""}
        ]
        with patch.object(
            ds, "_load_from_url", return_value=fallback_data
        ) as mock_url:
            with patch.dict("sys.modules", {"datasets": None}):
                result = ds.load_data()
        mock_url.assert_called_once()
        assert result == fallback_data

    def test_import_error_falls_back_to_url(self):
        ds = _make_dataset()
        fallback_data = [
            {"id": "url_1", "problem": "Q", "answer": "A", "canary": ""}
        ]
        with patch.object(
            ds, "_load_from_url", return_value=fallback_data
        ) as mock_url:
            with patch.dict("sys.modules", {"datasets": None}):
                result = ds.load_data()
        mock_url.assert_called_once()
        assert result == fallback_data


# ---------------------------------------------------------------------------
# load_data() – datasets exception falls back to _load_from_url
# ---------------------------------------------------------------------------


class TestLoadDatasetsExceptionFallback:
    def test_load_datasets_exception_fallback(self):
        """When load_dataset raises an exception, falls back to _load_from_url."""
        ds = _make_dataset()
        fallback_data = [
            {"id": "fb_1", "problem": "Q", "answer": "A", "canary": ""}
        ]
        mock_load_dataset = Mock(side_effect=RuntimeError("network error"))
        with patch.object(
            ds, "_load_from_url", return_value=fallback_data
        ) as mock_url:
            with patch.dict(
                "sys.modules",
                {"datasets": Mock(load_dataset=mock_load_dataset)},
            ):
                with patch(f"{MODULE}.logger"):
                    result = ds.load_data()
        mock_url.assert_called_once()
        assert result == fallback_data

    def test_exception_in_load_dataset_falls_back_to_url(self):
        ds = _make_dataset()
        fallback_data = [
            {"id": "fb_1", "problem": "Q", "answer": "A", "canary": ""}
        ]
        mock_load_dataset = Mock(side_effect=RuntimeError("network error"))
        with patch.object(
            ds, "_load_from_url", return_value=fallback_data
        ) as mock_url:
            with patch.dict(
                "sys.modules",
                {"datasets": Mock(load_dataset=mock_load_dataset)},
            ):
                with patch(f"{MODULE}.logger"):
                    result = ds.load_data()
        mock_url.assert_called_once()
        assert result == fallback_data


# ---------------------------------------------------------------------------
# _load_from_url – success and failure paths
# ---------------------------------------------------------------------------


class TestLoadFromUrl:
    def test_load_from_url_success(self):
        """_load_from_url calls pd.read_parquet and returns a list of dicts."""
        ds = _make_dataset()
        mock_df = MagicMock()
        row = MagicMock()
        row.get = lambda key, default="": {
            "id": "url_q1",
            "prompt": "Direct URL question",
            "answer": "Direct URL answer",
            "canary": "",
            "reference_steps": "",
        }.get(key, 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()
        mock_pd.read_parquet.assert_called_once()
        assert len(result) == 1
        assert result[0]["problem"] == "Direct URL question"

    def test_load_from_url_success_returns_list(self):
        ds = _make_dataset()
        mock_df = MagicMock()
        row = MagicMock()
        row.get = lambda key, default="": {
            "id": "url_q1",
            "prompt": "Direct URL question",
            "answer": "Direct URL answer",
            "canary": "",
            "reference_steps": "",
        }.get(key, 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"] == "Direct URL question"

    def test_load_from_url_failure_returns_empty(self):
        """_load_from_url returns [] when an exception occurs (e.g. connection error)."""
        ds = _make_dataset()
        mock_pd = Mock()
        mock_pd.read_parquet.side_effect = Exception("connection refused")
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger"):
                result = ds._load_from_url()
        assert result == []

    def test_load_from_url_exception_returns_empty_list(self):
        ds = _make_dataset()
        mock_pd = Mock()
        mock_pd.read_parquet.side_effect = Exception("connection refused")
        with patch.dict("sys.modules", {"pandas": mock_pd}):
            with patch(f"{MODULE}.logger"):
                result = ds._load_from_url()
        assert result == []