haiku.rag/evaluations/tests/test_datasets.py

from pathlib import Path

from evaluations.datasets.open_rag_bench import (
    build_orb_case,
    download_pdf,
    is_multimodal_query,
    map_orb_document,
    map_orb_retrieval,
)
from evaluations.datasets.wix import (
    build_wix_case,
    map_wix_document,
    map_wix_retrieval,
)


class TestWix:
    def test_map_document_with_all_fields(self) -> None:
        doc = {
            "id": 123,
            "url": "https://wix.com/article",
            "html_content": "<p>Content</p>",
            "title": "My Article",
        }
        payload = map_wix_document(doc)
        assert payload.uri == "123"
        assert payload.content == "<p>Content</p>"
        assert payload.title == "My Article"
        assert payload.format == "html"
        assert payload.metadata == {
            "article_id": "123",
            "url": "https://wix.com/article",
        }

    def test_map_document_no_id(self) -> None:
        doc = {
            "id": None,
            "url": "https://wix.com/page",
            "html_content": "<p>Text</p>",
            "title": None,
        }
        payload = map_wix_document(doc)
        assert payload.uri == "https://wix.com/page"

    def test_map_document_no_metadata(self) -> None:
        doc = {"id": None, "url": None, "html_content": "<p>X</p>", "title": None}
        payload = map_wix_document(doc)
        assert payload.metadata is None

    def test_map_retrieval(self) -> None:
        doc = {"question": "How to add a page?", "article_ids": [10, 20]}
        sample = map_wix_retrieval(doc)
        assert sample is not None
        assert sample.question == "How to add a page?"
        assert sample.expected_uris == ("10", "20")

    def test_map_retrieval_no_article_ids(self) -> None:
        doc = {"question": "Q?", "article_ids": None}
        assert map_wix_retrieval(doc) is None

    def test_map_retrieval_empty_article_ids(self) -> None:
        doc = {"question": "Q?", "article_ids": []}
        assert map_wix_retrieval(doc) is None

    def test_build_case(self) -> None:
        doc = {
            "question": "How?",
            "answer": "Like this.",
            "article_ids": [5, 10],
        }
        case = build_wix_case(2, doc)
        assert case.name == "2_5-10"
        assert case.inputs == "How?"
        assert case.expected_output == "Like this."
        assert case.metadata is not None
        assert case.metadata["case_index"] == "2"

    def test_build_case_no_article_ids(self) -> None:
        doc = {"question": "Q?", "answer": "A.", "article_ids": None}
        case = build_wix_case(1, doc)
        assert case.name == "case_1"


class TestOpenRAGBench:
    def test_map_document(self, tmp_path: Path) -> None:
        # Pre-create a cached PDF
        cache_dir = tmp_path / "pdfs"
        cache_dir.mkdir()
        pdf_path = cache_dir / "paper1.pdf"
        pdf_path.write_bytes(b"%PDF-fake")

        doc = {"paper_id": "paper1", "pdf_url": "https://example.com/paper1.pdf"}
        # Patch get_cache_dir to use our tmp_path
        from unittest.mock import patch

        with patch(
            "evaluations.datasets.open_rag_bench.get_cache_dir", return_value=cache_dir
        ):
            payload = map_orb_document(doc)

        assert payload is not None
        assert payload.uri == "paper1"
        assert payload.title == "paper1"
        assert payload.source_path == pdf_path
        assert payload.metadata == {"arxiv_id": "paper1"}

    def test_map_document_download_fails(self, tmp_path: Path) -> None:
        cache_dir = tmp_path / "pdfs"
        cache_dir.mkdir()

        doc = {"paper_id": "missing", "pdf_url": "https://example.com/missing.pdf"}
        from unittest.mock import patch

        with patch(
            "evaluations.datasets.open_rag_bench.get_cache_dir", return_value=cache_dir
        ):
            with patch(
                "evaluations.datasets.open_rag_bench.download_pdf", return_value=None
            ):
                payload = map_orb_document(doc)

        assert payload is None

    def test_map_retrieval(self) -> None:
        doc = {
            "query": "What is attention?",
            "doc_id": "1706.03762",
            "source": "text",
        }
        sample = map_orb_retrieval(doc)
        assert sample is not None
        assert sample.question == "What is attention?"
        assert sample.expected_uris == ("1706.03762",)
        assert sample.source_type == "text"

    def test_build_case(self) -> None:
        doc = {
            "query_id": "q_abcdef12",
            "query": "Explain transformers.",
            "answer": "Transformers are...",
            "type": "factual",
            "source": "text",
        }
        case = build_orb_case(1, doc)
        assert case.name == "1_q_abcdef"
        assert case.inputs == "Explain transformers."
        assert case.expected_output == "Transformers are..."
        assert case.metadata is not None
        assert case.metadata["query_id"] == "q_abcdef12"

    def test_download_pdf_uses_cache(self, tmp_path: Path) -> None:
        pdf_path = tmp_path / "cached.pdf"
        pdf_path.write_bytes(b"%PDF-cached")
        result = download_pdf("cached", "https://example.com/cached.pdf", tmp_path)
        assert result == pdf_path

    def test_is_multimodal_query(self) -> None:
        assert is_multimodal_query("image") is True
        assert is_multimodal_query("image_table") is True
        assert is_multimodal_query("text") is False