RAG-CUT/backend/tests/test_mineru_adapter.py

"""Tests for MinerU content_list → Block mapping."""

from __future__ import annotations

import tempfile
import subprocess
import unittest
from pathlib import Path
from unittest.mock import Mock, patch

from rag_cut.models import BlockType
from rag_cut.parsers.mineru_adapter import (
    _blocks_from_content_list,
    _find_content_list,
    _run_mineru,
)


class MineruAdapterProcessTest(unittest.TestCase):
    def test_default_timeout_allows_large_manual_to_finish(self) -> None:
        with patch.dict("os.environ", {}, clear=True):
            from rag_cut.parsers.mineru_adapter import _mineru_timeout_sec

            self.assertEqual(_mineru_timeout_sec(), 540.0)

    def test_configured_api_url_is_passed_to_mineru(self) -> None:
        process = Mock()
        process.communicate.return_value = ("", "")
        process.returncode = 0

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch("rag_cut.parsers.mineru_adapter._mineru_command", return_value="mineru"),
            patch.dict("os.environ", {"RAG_CUT_MINERU_API_URL": "http://127.0.0.1:30000"}),
            patch("rag_cut.parsers.mineru_adapter.subprocess.Popen", return_value=process) as popen,
        ):
            ok = _run_mineru(Path(tmp) / "manual.pdf", Path(tmp) / "output")

        self.assertTrue(ok)
        args = popen.call_args.args[0]
        self.assertEqual(args[-2:], ["--api-url", "http://127.0.0.1:30000"])

    def test_timeout_terminates_windows_process_tree(self) -> None:
        process = Mock()
        process.pid = 4321
        process.communicate.side_effect = subprocess.TimeoutExpired("mineru", 30)
        process.poll.return_value = None

        with (
            tempfile.TemporaryDirectory() as tmp,
            patch("rag_cut.parsers.mineru_adapter._mineru_command", return_value="mineru"),
            patch("rag_cut.parsers.mineru_adapter._mineru_timeout_sec", return_value=30),
            patch("rag_cut.parsers.mineru_adapter.os.name", "nt"),
            patch("rag_cut.parsers.mineru_adapter.subprocess.Popen", return_value=process),
            patch("rag_cut.parsers.mineru_adapter.subprocess.run") as run,
        ):
            ok = _run_mineru(Path(tmp) / "manual.pdf", Path(tmp) / "output")

        self.assertFalse(ok)
        run.assert_called_once_with(
            ["taskkill", "/PID", "4321", "/T", "/F"],
            check=False,
            capture_output=True,
            timeout=10,
        )


class MineruAdapterMappingTest(unittest.TestCase):
    def test_chart_becomes_image_block(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            assets = Path(tmp)
            img = assets / "chart.jpg"
            img.write_bytes(b"fake")
            items = [
                {
                    "type": "text",
                    "text": "第二十節 TradingView",
                    "text_level": 2,
                    "page_idx": 30,
                    "bbox": [127, 86, 363, 104],
                },
                {
                    "type": "chart",
                    "img_path": "chart.jpg",
                    "page_idx": 30,
                    "bbox": [129, 171, 878, 428],
                    "image_caption": [],
                },
            ]
            with patch("rag_cut.parsers.mineru_adapter._mineru_ocr_enabled", return_value=False):
                blocks = _blocks_from_content_list(items, assets / "content_list.json", assets)

            types = [b.type for b in blocks]
            self.assertEqual(types, [BlockType.HEADING, BlockType.IMAGE])
            self.assertEqual(blocks[1].image_id, "chart.jpg")
            self.assertEqual(blocks[1].meta.get("mineru_type"), "chart")
            self.assertEqual(blocks[1].meta.get("page"), 31)

    def test_table_keeps_markdown_and_screenshot(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            assets = Path(tmp)
            img = assets / "table.png"
            img.write_bytes(b"png")
            items = [
                {
                    "type": "table",
                    "img_path": "table.png",
                    "page_idx": 1,
                    "bbox": [100, 100, 400, 300],
                    "table_body": "<table><tr><td>A</td><td>B</td></tr></table>",
                }
            ]
            with patch("rag_cut.parsers.mineru_adapter._mineru_ocr_enabled", return_value=False):
                blocks = _blocks_from_content_list(items, assets / "content_list.json", assets)

            self.assertEqual(len(blocks), 1)
            self.assertEqual(blocks[0].type, BlockType.TABLE)
            self.assertIn("<table>", blocks[0].markdown or "")
            self.assertEqual(blocks[0].image_id, "table.png")
            self.assertTrue(blocks[0].image_path)

    def test_page_footnote_kept_as_paragraph(self) -> None:
        items = [
            {
                "type": "page_footnote",
                "text": "<sup>1</sup> See Market Master.",
                "page_idx": 8,
                "bbox": [85, 889, 912, 916],
            }
        ]
        with tempfile.TemporaryDirectory() as tmp:
            blocks = _blocks_from_content_list(items, Path(tmp) / "content_list.json", Path(tmp))
        self.assertEqual(len(blocks), 1)
        self.assertEqual(blocks[0].type, BlockType.PARAGRAPH)
        self.assertTrue(blocks[0].meta.get("is_footnote"))

    def test_list_items_become_paragraph(self) -> None:
        items = [
            {
                "type": "list",
                "list_items": ["第一點說明", {"text": "第二點說明"}],
                "page_idx": 2,
                "bbox": [100, 200, 400, 260],
            }
        ]
        with tempfile.TemporaryDirectory() as tmp:
            blocks = _blocks_from_content_list(items, Path(tmp) / "content_list.json", Path(tmp))
        self.assertEqual(len(blocks), 1)
        self.assertEqual(blocks[0].type, BlockType.PARAGRAPH)
        self.assertIn("第一點說明", blocks[0].text or "")
        self.assertIn("第二點說明", blocks[0].text or "")

    def test_find_content_list_prefers_non_v2(self) -> None:
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp) / "doc" / "auto"
            root.mkdir(parents=True)
            v1 = root / "doc_content_list.json"
            v2 = root / "doc_content_list_v2.json"
            v1.write_text("[]", encoding="utf-8")
            v2.write_text("[]", encoding="utf-8")
            found = _find_content_list(Path(tmp))
            self.assertEqual(found, v1)


if __name__ == "__main__":
    unittest.main()