DocsGPT/tests/parser/file/test_pdf_trust.py

"""Tests for the PDF trust checks behind ``PDF_TRUST_CHECK``.

The checks target the two classes where anydoc drops text *silently*:
composite (Type0) fonts without a ToUnicode map, and CJK-declaring PDFs
whose extracted text carries almost no CJK. The scanner is byte-level and
regex-based, so the synthetic fixtures below are minimal PDF fragments, not
well-formed files; the one real fixture is the corpus document that loses
its whole Chinese column in anydoc with no error.
"""
import zlib
from pathlib import Path

from application.parser.file.pdf_trust import (
    check_pdf_fonts,
    verify_extraction,
    verify_pdf_file,
)

FIXTURES = Path(__file__).parent / "fixtures"


def _obj(body: bytes) -> bytes:
    return b"1 0 obj" + body + b"endobj\n"


def _stream(payload: bytes) -> bytes:
    return b"2 0 obj<</Filter/FlateDecode>>stream\n" + zlib.compress(payload) + b"\nendstream endobj\n"


TYPE0_WITH_TOUNICODE = _obj(
    b"<</Type/Font/Subtype/Type0/BaseFont/AAAAAA+Simple/ToUnicode 9 0 R>>"
)
TYPE0_BARE = _obj(b"<</Type/Font/Subtype/Type0/BaseFont/AAAAAA+Bare>>")
SIMPLE_FONT = _obj(b"<</Type/Font/Subtype/Type1/BaseFont/Helvetica>>")
GB1_ORDERING = _obj(b"<</Registry(Adobe)/Ordering(GB1)/Supplement 5>>")


class TestCheckPdfFonts:
    def test_type0_with_tounicode_not_flagged(self):
        result = check_pdf_fonts(TYPE0_WITH_TOUNICODE)
        assert result["type0"] == 1
        assert result["type0_no_tounicode"] == 0
        assert not result["flagged"]

    def test_type0_without_tounicode_flagged(self):
        result = check_pdf_fonts(TYPE0_BARE)
        assert result["type0_no_tounicode"] == 1
        assert result["flagged"]

    def test_simple_font_not_flagged(self):
        result = check_pdf_fonts(SIMPLE_FONT)
        assert result["type0"] == 0
        assert result["has_fonts"]
        assert not result["flagged"]

    def test_no_fonts_flagged(self):
        result = check_pdf_fonts(_obj(b"<</Type/Page>>"))
        assert not result["has_fonts"]
        assert result["flagged"]

    def test_type0_inside_object_stream_is_seen(self):
        """Object streams hold dicts with no obj markers; the scan must decompress them."""
        data = _stream(b"<</Type/Font/Subtype/Type0/BaseFont/BBBBBB+Packed>>")
        result = check_pdf_fonts(data)
        assert result["type0"] == 1
        assert result["type0_no_tounicode"] == 1
        assert result["flagged"]

    def test_font_inside_object_stream_counts_as_has_fonts(self):
        """A fully-compressed PDF with only simple fonts must not false-positive."""
        data = _stream(b"<</Type/Font/Subtype/TrueType/BaseFont/Arial>>")
        result = check_pdf_fonts(data)
        assert result["has_fonts"]
        assert not result["flagged"]

    def test_cjk_ordering_detected_in_stream(self):
        data = _stream(b"<</Registry(Adobe)/Ordering(CNS1)/Supplement 4>>") + SIMPLE_FONT
        assert check_pdf_fonts(data)["expects_cjk"]

    def test_cjk_font_names_alone_do_not_set_expectation(self):
        """Regression for the Berkshire false positive: Word-exported Latin PDFs
        embed MS-Gothic subsets for stray full-width characters, and substring
        matching also hits FranklinGothic — neither means CJK *content*, and a
        false expectation reroutes a 150-page English report to the heavy
        engine. Only CIDSystemInfo /Ordering may set the expectation."""
        data = _obj(
            b"<</Type/Font/Subtype/TrueType/BaseFont/AKBHWD+MS-Gothic>>"
        ) + _obj(b"<</Type/Font/Subtype/TrueType/BaseFont/KMBGKH+FranklinGothic-Roman>>")
        result = check_pdf_fonts(data)
        assert not result["expects_cjk"]
        assert not result["flagged"]


class TestVerifyExtraction:
    def test_clean_pdf_and_output(self):
        assert verify_extraction(SIMPLE_FONT, "Plain English text.") == []

    def test_bare_type0_reported(self):
        problems = verify_extraction(TYPE0_BARE + SIMPLE_FONT, "some text")
        assert len(problems) == 1
        assert "ToUnicode" in problems[0]

    def test_cjk_expected_but_missing(self):
        problems = verify_extraction(GB1_ORDERING + SIMPLE_FONT, "English only output")
        assert len(problems) == 1
        assert "CJK" in problems[0]

    def test_cjk_expected_and_present(self):
        text = "标题：本协议由双方共同签署生效" # 14 CJK chars
        assert verify_extraction(GB1_ORDERING + SIMPLE_FONT, text) == []

    def test_both_problems_reported(self):
        problems = verify_extraction(TYPE0_BARE + GB1_ORDERING, "english")
        assert len(problems) == 2


class TestRealFixture:
    """The corpus PDF where anydoc silently drops the entire Chinese column."""

    def test_cid_font_nda_is_flagged(self):
        data = (FIXTURES / "nda_en_zh_cid_font.pdf").read_bytes()
        result = check_pdf_fonts(data)
        assert result["type0_no_tounicode"] >= 1
        assert result["expects_cjk"]
        assert result["flagged"]

    def test_cjkless_extraction_reports_both(self):
        data = (FIXTURES / "nda_en_zh_cid_font.pdf").read_bytes()
        problems = verify_extraction(data, "MUTUAL NON-DISCLOSURE AGREEMENT ...")
        assert any("ToUnicode" in p for p in problems)
        assert any("CJK" in p for p in problems)


def test_flate_bomb_streams_are_capped():
    """Flate reaches ~1000:1, so per-stream inflation must be capped — an
    uncapped decompress of a crafted PDF would OOM the ingest worker."""
    from application.parser.file.pdf_trust import _decompressed_streams, _STREAM_INFLATE_CAP

    bomb = _stream(b"\0" * (_STREAM_INFLATE_CAP * 4))
    chunks = list(_decompressed_streams(bomb))
    assert chunks
    assert all(len(chunk) <= _STREAM_INFLATE_CAP for chunk in chunks)
    # The capped scan still completes and reports sanely alongside real objects.
    assert not check_pdf_fonts(SIMPLE_FONT + bomb)["flagged"]


def test_verify_pdf_file_unreadable_trusts_output(tmp_path):
    assert verify_pdf_file(tmp_path / "missing.pdf", "text") == []