mcp-nextcloud/tests/unit/test_doc_classifier.py

"""Unit tests for the tier-0 document classifier.

Pins the routing decisions and the text-quality heuristic that drive which
extraction tier a PDF starts in:
  * a clean born-digital PDF (text, no full-page images) -> ``fast`` (tier 1);
  * a full-page-image scan -> ``ocr`` (tier 3), since handwriting/stamps aren't
    in any text layer;
  * the text-quality score distinguishes clean prose from mashed/space-less junk.
"""

import pymupdf
import pytest

from nextcloud_mcp_server.document_processors import classifier as clf

pytestmark = pytest.mark.unit


def _digital_pdf(
    pages: int = 3, body: str = "Hello world this is clean text. "
) -> bytes:
    doc = pymupdf.open()
    for _ in range(pages):
        page = doc.new_page(width=595, height=842)
        page.insert_text((50, 60), body * 8)
    data: bytes = doc.tobytes()
    doc.close()
    return data


def _full_page_image_pdf(pages: int = 2) -> bytes:
    # A page whose entire area is a raster image -> looks scanned.
    doc = pymupdf.open()
    pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 600, 850))
    pix.clear_with(255)
    img = pix.tobytes("png")
    del pix  # Pixmap holds native memory; release it before the loop
    for _ in range(pages):
        page = doc.new_page(width=595, height=842)
        page.insert_image(page.rect, stream=img)
    data: bytes = doc.tobytes()
    doc.close()
    return data


# --- text-quality heuristic --------------------------------------------------


def test_text_quality_clean_prose_scores_high():
    assert clf._text_quality("the quick brown fox jumps over the lazy dog") > 0.8


def test_text_quality_mashed_tokens_scores_low():
    # space-less / mashed layer (the "Student 147" failure mode)
    mashed = "01322234567mobileoutstandingresilienceacademicachievementhurdles"
    assert clf._text_quality(mashed) < clf.MIN_TEXT_QUALITY


def test_text_quality_empty_is_zero():
    assert clf._text_quality("") == pytest.approx(0.0)


# --- routing -----------------------------------------------------------------


def test_digital_pdf_routes_fast():
    c = clf.classify_pdf(_digital_pdf())
    assert c.recommended_tier == "fast"
    assert c.ocr_page_fraction == pytest.approx(0.0)
    assert "image_heavy" not in c.flags
    assert c.mean_text_quality > 0.8


def test_full_page_image_routes_ocr():
    c = clf.classify_pdf(_full_page_image_pdf())
    assert c.recommended_tier == "ocr"
    assert c.ocr_page_fraction == pytest.approx(1.0)
    assert "image_heavy" in c.flags
    assert "scanned" in c.flags  # no text layer at all


# --- sampling bounds large docs ----------------------------------------------


def test_large_doc_is_sampled():
    c = clf.classify_pdf(_digital_pdf(pages=120))
    assert c.page_count == 120
    assert c.sampled_pages <= clf.MAX_SAMPLED_PAGES


def test_sample_indices_includes_first_and_last_page():
    idx = clf._sample_indices(100)
    assert idx[0] == 0
    assert idx[-1] == 99  # last page must be sampled (scanned-tail case)
    assert len(idx) <= clf.MAX_SAMPLED_PAGES


# --- flag paths --------------------------------------------------------------


def _image_with_mashed_text_pdf(pages: int = 2) -> bytes:
    # Full-page image with a junk (mashed/space-less) text layer over it -- a
    # scan whose OCR'd text layer is unusable.
    doc = pymupdf.open()
    pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 600, 850))
    pix.clear_with(255)
    img = pix.tobytes("png")
    del pix  # Pixmap holds native memory; release it before the loop
    mashed = "01322234567mobileoutstandingresilienceacademicachievement " * 3
    for _ in range(pages):
        page = doc.new_page(width=595, height=842)
        page.insert_image(page.rect, stream=img)
        page.insert_text((50, 60), mashed)
    data: bytes = doc.tobytes()
    doc.close()
    return data


def test_scanned_flag_when_no_text_layer():
    c = clf.classify_pdf(_full_page_image_pdf())
    assert c.total_chars == 0
    assert "scanned" in c.flags
    assert c.recommended_tier == "ocr"


def test_bad_text_layer_flag_on_image_with_junk_text():
    c = clf.classify_pdf(_image_with_mashed_text_pdf())
    assert c.total_chars > 0
    assert c.mean_text_quality < clf.MIN_TEXT_QUALITY
    assert "bad_text_layer" in c.flags
    assert c.recommended_tier == "ocr"


def _mostly_text_one_image_pdf() -> bytes:
    # 3 digital text pages + 1 full-page-image page: one image-heavy page, but
    # ocr_frac = 1/4 < OCR_PAGE_FRACTION, so the doc routes fast.
    doc = pymupdf.open()
    pix = pymupdf.Pixmap(pymupdf.csRGB, pymupdf.IRect(0, 0, 600, 850))
    pix.clear_with(255)
    img = pix.tobytes("png")
    del pix  # Pixmap holds native memory; release it before the loop
    for _ in range(3):
        page = doc.new_page(width=595, height=842)
        page.insert_text((50, 60), "Hello world this is clean text. " * 8)
    page = doc.new_page(width=595, height=842)
    page.insert_image(page.rect, stream=img)
    data: bytes = doc.tobytes()
    doc.close()
    return data


def test_image_heavy_flag_without_ocr_routing():
    # The documented asymmetry operators rely on: a mostly-digital doc with one
    # full-page image carries the image_heavy flag yet still routes fast.
    c = clf.classify_pdf(_mostly_text_one_image_pdf())
    assert "image_heavy" in c.flags
    assert c.recommended_tier == "fast"
    assert c.ocr_page_fraction < clf.OCR_PAGE_FRACTION


# --- classify_from_text (hot-path, derived from tier-1 extraction) -----------


def test_classify_from_text_clean_routes_fast():
    txt = "the quick brown fox jumps over the lazy dog " * 3
    c = clf.classify_from_text(
        txt, [{"page": 1, "start_offset": 0, "end_offset": len(txt)}]
    )
    assert c.recommended_tier == "fast"
    assert c.mean_text_quality > 0.8
    assert c.flags == set()


def test_classify_from_text_empty_routes_ocr():
    c = clf.classify_from_text("", [{"page": 1, "start_offset": 0, "end_offset": 0}])
    assert c.recommended_tier == "ocr"
    assert "no_text_layer" in c.flags
    assert c.total_chars == 0