Spaces:
Runtime error
Runtime error
| from __future__ import annotations | |
| from pathlib import Path | |
| from typing import TYPE_CHECKING | |
| from unittest.mock import MagicMock, PropertyMock | |
| import pytest | |
| from kotaemon.base import Document | |
| from kotaemon.loaders.paddleocr_loader.adapter import PaddleOCRResult | |
| from .conftest import skip_when_paddleocr_not_installed | |
| if TYPE_CHECKING: | |
| from pytest_mock.plugin import MockerFixture | |
| def _make_page_result(inner_dict: dict) -> MagicMock: | |
| """Build a mock page result with .json['res'] as used by the adapter.""" | |
| page = MagicMock() | |
| page.json = {"res": inner_dict} | |
| return page | |
| # --------------------------------------------------------------------------- | |
| # PaddleOCRResult adapter tests | |
| # --------------------------------------------------------------------------- | |
| def test_clean_table_html_strips_wrapper() -> None: | |
| """_clean_table_html removes <html><body> wrapper.""" | |
| raw = [_make_page_result({"page_index": 0, "parsing_res_list": []})] | |
| result = PaddleOCRResult( | |
| raw_result=raw, file_path=Path("/tmp/doc.pdf"), extra_info={} | |
| ) | |
| out = result._clean_table_html("<html><body><table>x</table></body></html>") | |
| assert out.strip() == "<table>x</table>" | |
| def test_file_name_property() -> None: | |
| """file_name returns the path name.""" | |
| raw = [_make_page_result({"page_index": 0, "parsing_res_list": []})] | |
| result = PaddleOCRResult( | |
| raw_result=raw, | |
| file_path=Path("/some/dir/doc.pdf"), | |
| extra_info={}, | |
| ) | |
| assert result.file_name == "doc.pdf" | |
| def test_to_documents_text_and_table() -> None: | |
| """to_documents returns text and table docs from parsing_res_list.""" | |
| raw = [ | |
| _make_page_result( | |
| { | |
| "page_index": 0, | |
| "parsing_res_list": [ | |
| {"block_label": "paragraph_title", "block_content": "Title"}, | |
| {"block_label": "text", "block_content": "Body text."}, | |
| { | |
| "block_label": "table", | |
| "block_content": "<html><body><table><tr><td>A</td></tr>" | |
| "</table></body></html>", | |
| }, | |
| ], | |
| } | |
| ) | |
| ] | |
| result = PaddleOCRResult( | |
| raw_result=raw, | |
| file_path=Path("doc.pdf"), | |
| extra_info={"source": "test"}, | |
| ) | |
| docs = result.to_documents() | |
| text_docs = [d for d in docs if d.metadata.get("type") != "table"] | |
| table_docs = [d for d in docs if d.metadata.get("type") == "table"] | |
| assert len(text_docs) == 1 | |
| assert "Title" in text_docs[0].text and "Body text." in text_docs[0].text | |
| assert text_docs[0].metadata.get("source") == "test" | |
| assert len(table_docs) == 1 | |
| assert "A" in table_docs[0].text | |
| assert table_docs[0].metadata["table_origin"] | |
| assert table_docs[0].metadata["page_label"] == 1 | |
| def test_to_documents_skips_empty_content() -> None: | |
| """Blocks with empty block_content are skipped.""" | |
| raw = [ | |
| _make_page_result( | |
| { | |
| "page_index": 0, | |
| "parsing_res_list": [ | |
| {"block_label": "text", "block_content": ""}, | |
| {"block_label": "paragraph_title", "block_content": "Only"}, | |
| ], | |
| } | |
| ) | |
| ] | |
| result = PaddleOCRResult(raw_result=raw, file_path=Path("x.pdf"), extra_info={}) | |
| docs = result.to_documents() | |
| assert len(docs) == 1 | |
| assert docs[0].text.strip() == "Only" | |
| # --------------------------------------------------------------------------- | |
| # PPStructureV3Reader tests | |
| # --------------------------------------------------------------------------- | |
| def test_ppstructure_v3_unsupported_file_type_raises() -> None: | |
| """load_data raises ValueError for unsupported extension.""" | |
| from kotaemon.loaders import PPStructureV3Reader | |
| reader = PPStructureV3Reader() | |
| with pytest.raises(ValueError, match="Unsupported file type"): | |
| reader.load_data(Path("/tmp/file.xyz")) | |
| def test_ppstructure_v3_supported_file_types_attribute() -> None: | |
| """Reader has expected supported_file_types.""" | |
| from kotaemon.loaders import PPStructureV3Reader | |
| reader = PPStructureV3Reader() | |
| assert ".pdf" in reader.supported_file_types | |
| assert ".png" in reader.supported_file_types | |
| def test_ppstructure_v3_load_data_with_mocked_pipeline( | |
| mocker: MockerFixture, | |
| ) -> None: | |
| """load_data returns list of Documents when pipeline is mocked.""" | |
| from kotaemon.loaders.paddleocr_loader.ppstructure_v3_loader import ( | |
| PPStructureV3Reader, | |
| ) | |
| mock_raw = [ | |
| _make_page_result( | |
| { | |
| "page_index": 0, | |
| "parsing_res_list": [ | |
| {"block_label": "text", "block_content": "Hello"}, | |
| ], | |
| } | |
| ) | |
| ] | |
| mock_pipeline = MagicMock() | |
| mock_pipeline.predict.return_value = mock_raw | |
| mocker.patch.object( | |
| PPStructureV3Reader, | |
| "pipeline_", | |
| new_callable=PropertyMock, | |
| return_value=mock_pipeline, | |
| ) | |
| reader = PPStructureV3Reader() | |
| docs = reader.load_data(Path("/tmp/doc.pdf"), extra_info={}) | |
| assert len(docs) >= 1 | |
| assert all(isinstance(d, Document) for d in docs) | |
| mock_pipeline.predict.assert_called_once_with("/tmp/doc.pdf") | |
| def test_ppstructure_v3_cpu_disables_mkldnn( | |
| mocker: MockerFixture, | |
| device: str, | |
| ) -> None: | |
| """CPU init passes enable_mkldnn=False (Paddle#77340 workaround).""" | |
| from kotaemon.loaders.paddleocr_loader import ppstructure_v3_loader | |
| mock_pipeline = MagicMock() | |
| mock_ctor = mocker.patch( | |
| "paddleocr.PPStructureV3", | |
| return_value=mock_pipeline, | |
| ) | |
| reader = ppstructure_v3_loader.PPStructureV3Reader(device=device) | |
| _ = reader.pipeline_ | |
| mock_ctor.assert_called_once() | |
| assert mock_ctor.call_args.kwargs["enable_mkldnn"] is False | |
| def test_ppstructure_v3_gpu_omits_mkldnn_override( | |
| mocker: MockerFixture, | |
| ) -> None: | |
| """GPU init leaves enable_mkldnn to paddleocr defaults.""" | |
| from kotaemon.loaders.paddleocr_loader import ppstructure_v3_loader | |
| mock_pipeline = MagicMock() | |
| mock_ctor = mocker.patch( | |
| "paddleocr.PPStructureV3", | |
| return_value=mock_pipeline, | |
| ) | |
| reader = ppstructure_v3_loader.PPStructureV3Reader(device="gpu:0") | |
| _ = reader.pipeline_ | |
| mock_ctor.assert_called_once() | |
| assert "enable_mkldnn" not in mock_ctor.call_args.kwargs | |
| # --------------------------------------------------------------------------- | |
| # PaddleOCRVLReader tests | |
| # --------------------------------------------------------------------------- | |
| def test_paddleocr_vl_unsupported_file_type_raises() -> None: | |
| """load_data raises ValueError for unsupported extension.""" | |
| from kotaemon.loaders import PaddleOCRVLReader | |
| reader = PaddleOCRVLReader() | |
| with pytest.raises(ValueError, match="Unsupported file type"): | |
| reader.load_data(Path("/tmp/file.xyz")) | |
| def test_paddleocr_vl_supported_file_types_attribute() -> None: | |
| """Reader has expected supported_file_types.""" | |
| from kotaemon.loaders import PaddleOCRVLReader | |
| reader = PaddleOCRVLReader() | |
| assert ".pdf" in reader.supported_file_types | |
| assert ".png" in reader.supported_file_types | |
| def test_paddleocr_vl_load_data_with_mocked_pipeline( | |
| mocker: MockerFixture, | |
| ) -> None: | |
| """load_data returns list of Documents when pipeline is mocked.""" | |
| from kotaemon.loaders import PaddleOCRVLReader | |
| mock_raw = [ | |
| _make_page_result( | |
| { | |
| "page_index": 0, | |
| "parsing_res_list": [ | |
| {"block_label": "text", "block_content": "VL text"}, | |
| ], | |
| } | |
| ) | |
| ] | |
| mock_pipeline = MagicMock() | |
| mock_pipeline.predict.return_value = mock_raw | |
| mocker.patch.object( | |
| PaddleOCRVLReader, | |
| "pipeline_", | |
| new_callable=PropertyMock, | |
| return_value=mock_pipeline, | |
| ) | |
| reader = PaddleOCRVLReader() | |
| docs = reader.load_data(Path("/tmp/img.png")) | |
| assert len(docs) >= 1 | |
| assert all(isinstance(d, Document) for d in docs) | |
| mock_pipeline.predict.assert_called_once_with("/tmp/img.png") | |
| def test_paddleocr_vl_cpu_disables_mkldnn( | |
| mocker: MockerFixture, | |
| device: str, | |
| ) -> None: | |
| """CPU init passes enable_mkldnn=False (Paddle#77340 workaround).""" | |
| from kotaemon.loaders.paddleocr_loader import paddleocr_vl_loader | |
| mock_pipeline = MagicMock() | |
| mock_ctor = mocker.patch( | |
| "paddleocr.PaddleOCRVL", | |
| return_value=mock_pipeline, | |
| ) | |
| reader = paddleocr_vl_loader.PaddleOCRVLReader(device=device) | |
| _ = reader.pipeline_ | |
| mock_ctor.assert_called_once() | |
| assert mock_ctor.call_args.kwargs["enable_mkldnn"] is False | |
| def test_paddleocr_vl_gpu_omits_mkldnn_override( | |
| mocker: MockerFixture, | |
| ) -> None: | |
| """GPU init leaves enable_mkldnn to paddleocr defaults.""" | |
| from kotaemon.loaders.paddleocr_loader import paddleocr_vl_loader | |
| mock_pipeline = MagicMock() | |
| mock_ctor = mocker.patch( | |
| "paddleocr.PaddleOCRVL", | |
| return_value=mock_pipeline, | |
| ) | |
| reader = paddleocr_vl_loader.PaddleOCRVLReader(device="gpu:0") | |
| _ = reader.pipeline_ | |
| mock_ctor.assert_called_once() | |
| assert "enable_mkldnn" not in mock_ctor.call_args.kwargs | |