/tests/benchmarks/datasets/test_browsecomp_coverage.py

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

Targets the 37 missing lines:
- get_dataset_info / get_default_dataset_path
- process_example() with canary and without canary
- process_example() known_answer_map lookup
- process_example() still-encrypted answer fallback keys
- process_example() decryption exception
- get_question / get_answer
"""

from unittest.mock import patch

MODULE = "local_deep_research.benchmarks.datasets.browsecomp"


def _make_dataset(**kwargs):
    from local_deep_research.benchmarks.datasets.browsecomp import (
        BrowseCompDataset,
    )

    return BrowseCompDataset(**kwargs)


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


class TestDatasetInfo:
    def test_get_dataset_info_has_browsecomp_id(self):
        from local_deep_research.benchmarks.datasets.browsecomp import (
            BrowseCompDataset,
        )

        info = BrowseCompDataset.get_dataset_info()
        assert info["id"] == "browsecomp"
        assert "name" in info
        assert "description" in info

    def test_get_default_dataset_path_is_url(self):
        from local_deep_research.benchmarks.datasets.browsecomp import (
            BrowseCompDataset,
        )

        path = BrowseCompDataset.get_default_dataset_path()
        assert path.startswith("https://")
        assert "browse_comp" in path or "browsecomp" in path.lower()


# ---------------------------------------------------------------------------
# process_example – no canary
# ---------------------------------------------------------------------------


class TestProcessExampleNoCanary:
    @patch(f"{MODULE}.get_known_answer_map", return_value={})
    @patch(f"{MODULE}.decrypt")
    def test_no_canary_returns_original(self, mock_decrypt, mock_known_map):
        ds = _make_dataset()
        example = {"problem": "encoded_problem", "answer": "encoded_answer"}
        result = ds.process_example(example)
        # Without canary, decryption is skipped
        mock_decrypt.assert_not_called()
        assert result["problem"] == "encoded_problem"

    @patch(f"{MODULE}.get_known_answer_map", return_value={})
    @patch(f"{MODULE}.decrypt")
    def test_empty_canary_returns_original(self, mock_decrypt, mock_known_map):
        ds = _make_dataset()
        example = {
            "problem": "some_problem",
            "answer": "some_answer",
            "canary": "",
        }
        ds.process_example(example)
        mock_decrypt.assert_not_called()


# ---------------------------------------------------------------------------
# process_example – with canary, problem too short to decrypt
# ---------------------------------------------------------------------------


class TestProcessExampleShortProblem:
    @patch(f"{MODULE}.get_known_answer_map", return_value={})
    @patch(f"{MODULE}.decrypt", return_value="decrypted_answer")
    def test_short_problem_not_decrypted(self, mock_decrypt, mock_known_map):
        ds = _make_dataset()
        example = {
            "problem": "short",  # len <= 20, not decrypted
            "answer": "some_answer",
            "canary": "mykey",
        }
        ds.process_example(example)
        # decrypt called for answer but not problem
        decrypt_calls = [c[0][1] for c in mock_decrypt.call_args_list]
        assert "short" not in decrypt_calls


# ---------------------------------------------------------------------------
# process_example – known_answer_map lookup
# ---------------------------------------------------------------------------


class TestProcessExampleKnownAnswerMap:
    @patch(f"{MODULE}.get_known_answer_map")
    @patch(f"{MODULE}.decrypt")
    def test_known_answer_used_without_decrypt(
        self, mock_decrypt, mock_known_map
    ):
        known_map = {"encoded_answer_xyz": "Real Answer"}
        mock_known_map.return_value = known_map
        # decrypt returns something different to confirm it's not called for known answers
        mock_decrypt.return_value = "decrypted_problem_text"
        ds = _make_dataset()
        problem_str = "A" * 25  # long enough to trigger problem decryption
        example = {
            "problem": problem_str,
            "answer": "encoded_answer_xyz",
            "canary": "key",
        }
        result = ds.process_example(example)
        assert result["answer"] == "Real Answer"
        assert result["correct_answer"] == "Real Answer"


# ---------------------------------------------------------------------------
# process_example – regular decryption when not in known map
# ---------------------------------------------------------------------------


class TestProcessExampleRegularDecryption:
    @patch(f"{MODULE}.get_known_answer_map", return_value={})
    @patch(f"{MODULE}.decrypt")
    def test_decrypt_called_for_long_problem_and_answer(
        self, mock_decrypt, mock_known_map
    ):
        mock_decrypt.side_effect = lambda text, key: f"decrypted_{text}"
        ds = _make_dataset()
        long_problem = "P" * 25
        example = {
            "problem": long_problem,
            "answer": "some_answer",
            "canary": "key",
        }
        result = ds.process_example(example)
        assert result["problem"] == f"decrypted_{long_problem}"
        assert result["answer"] == "decrypted_some_answer"
        assert result["correct_answer"] == "decrypted_some_answer"

    @patch(f"{MODULE}.get_known_answer_map", return_value={})
    @patch(f"{MODULE}.decrypt")
    def test_no_answer_field_not_processed(self, mock_decrypt, mock_known_map):
        mock_decrypt.return_value = "decrypted"
        ds = _make_dataset()
        example = {"problem": "P" * 25, "canary": "key"}
        result = ds.process_example(example)
        # No answer field → answer decryption not attempted
        assert "answer" not in result or result.get("answer", "") == ""


# ---------------------------------------------------------------------------
# process_example – still-encrypted answer with fallback keys
# ---------------------------------------------------------------------------


class TestProcessExampleFallbackKeys:
    @patch(f"{MODULE}.get_known_answer_map", return_value={})
    @patch(f"{MODULE}.decrypt")
    def test_fallback_key_used_when_still_looks_encrypted(
        self, mock_decrypt, mock_known_map
    ):
        # First call (with canary) returns same base64 text (still encrypted)
        # Second call (with fallback key "MHGGF2022!") returns plaintext
        base64_answer = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789+/="
        call_count = [0]

        def decrypt_side_effect(text, key):
            call_count[0] += 1
            if call_count[0] == 1:
                return text  # problem decryption
            if call_count[0] == 2:
                return text  # answer with canary – still looks encrypted
            if key == "MHGGF2022!":
                return "Plaintext Answer From Fallback"
            return text

        mock_decrypt.side_effect = decrypt_side_effect
        ds = _make_dataset()
        example = {
            "problem": "P" * 25,
            "answer": base64_answer,
            "canary": "wrongkey",
        }
        result = ds.process_example(example)
        # The answer should have been updated by the fallback decryption
        assert "answer" in result


# ---------------------------------------------------------------------------
# process_example – decryption exception
# ---------------------------------------------------------------------------


class TestProcessExampleDecryptException:
    @patch(f"{MODULE}.get_known_answer_map", return_value={})
    @patch(f"{MODULE}.decrypt", side_effect=Exception("decode error"))
    def test_exception_during_decrypt_logged_not_raised(
        self, mock_decrypt, mock_known_map
    ):
        ds = _make_dataset()
        example = {
            "problem": "P" * 25,
            "answer": "some_answer",
            "canary": "key",
        }
        # Should not raise
        result = ds.process_example(example)
        # Result should still be a dict
        assert isinstance(result, dict)


# ---------------------------------------------------------------------------
# get_question / get_answer
# ---------------------------------------------------------------------------


class TestGetQuestionAndAnswer:
    def test_get_question_returns_problem(self):
        ds = _make_dataset()
        example = {"problem": "What is 2+2?", "answer": "4"}
        assert ds.get_question(example) == "What is 2+2?"

    def test_get_question_returns_empty_when_missing(self):
        ds = _make_dataset()
        assert ds.get_question({}) == ""

    def test_get_answer_prefers_correct_answer(self):
        ds = _make_dataset()
        example = {"answer": "raw_answer", "correct_answer": "preferred_answer"}
        assert ds.get_answer(example) == "preferred_answer"

    def test_get_answer_falls_back_to_answer(self):
        ds = _make_dataset()
        example = {"answer": "fallback"}
        assert ds.get_answer(example) == "fallback"