From 4e6a13576d6ee372ec7e62d53a1ae695b61d4060 Mon Sep 17 00:00:00 2001 From: mountain Date: Thu, 9 Jul 2026 11:58:59 +0800 Subject: [PATCH] fix: CMYK image drop, empty-doc crash, page_index shadowing, sqlite hardening, flaky tests Addresses items 9-13 and a/b/c/f from the max-effort review of PR #272. - pdf.py: image colorspace check was `pix.n > 4`, which treats CMYK-without- alpha (n==4, same as RGBA) as not needing RGB conversion; pix.save() as .png then raises "unsupported colorspace", silently dropped by the surrounding except. Fixed to `pix.n - pix.alpha >= 4` (correctly converts CMYK, leaves RGBA untouched). - pipeline.py: detect_strategy([]) (an empty/whitespace-only source file) returned "content_based", routing into the PDF-oriented TOC-detection pipeline -- wasting a real LLM call before raising IndexingError. Empty node lists now route to level_based, whose build_tree_from_levels([]) returns an empty structure instantly with zero LLM calls. - page_index.py (shim): pageindex/__init__.py binds the canonical `page_index` function as the package attribute, but this file is ALSO a real submodule of the same name -- importing it anywhere (import machinery, unconditional) overwrites that attribute with the module object, breaking `from pageindex import page_index; page_index(x)` for the rest of the process. Made the shim module itself callable (delegates to the real function via a ModuleType subclass), so whichever object ends up in that slot is callable regardless of import order. - storage/sqlite.py: create_collection let a raw sqlite3.IntegrityError escape on a duplicate name (new CollectionAlreadyExistsError); the collections table's CHECK constraint only validated the name's first character (GLOB '*' is a wildcard, not a regex quantifier over the preceding class) -- fixed to validate the whole string, and SQLiteStorage now also validates in Python (it's a public StorageEngine usable directly, bypassing LocalBackend's own check). - tests/test_review_fixes_2.py: two tests used a ContentNode with no `level` set, so build_index took the content_based path and made real (retried, slow, and -- with a valid key -- billable) LLM calls instead of testing the text-stripping logic they claimed to. Mocked out _content_based_pipeline. - retrieve.py: _parse_pages/_get_pdf_page_content were independent copies of the canonical parse_pages/get_pdf_page_content that had already drifted (missing the p>=1 filter and 1000-page DoS cap) -- delegate to canonical now, so the legacy pageindex.get_page_content path can't silently regress again. - parser/markdown.py: a leading UTF-8 BOM broke first-header detection (not whitespace, .strip() doesn't remove it) -- decode utf-8-sig. Only backtick fences were recognized as code blocks, so a '#'-prefixed line inside a ~~~-fenced block (valid CommonMark) was misparsed as a heading -- recognize both fence styles. - run_pageindex.py: --if-thinning wasn't migrated to the bare-flag + legacy-yes/no convention the other four --if-add-* flags got; bare usage raised an argparse error and it never went through the shared coercion. - types.py: DocumentDetail's `structure` field was inside the class's total=False body, so TypedDict rules made it optional even though every backend always populates it. Split into a required base class. Adds regression tests for all of the above. Full suite: 244 passed, 2 skipped (one pre-existing, unrelated flaky cloud-streaming test). Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS --- pageindex/errors.py | 5 ++++ pageindex/index/pipeline.py | 7 +++++ pageindex/page_index.py | 25 ++++++++++++++++ pageindex/parser/markdown.py | 12 ++++++-- pageindex/parser/pdf.py | 8 +++++- pageindex/retrieve.py | 40 ++++++++++---------------- pageindex/storage/sqlite.py | 33 +++++++++++++++++++-- pageindex/types.py | 16 +++++++++-- run_pageindex.py | 6 ++-- tests/test_legacy_shims.py | 29 +++++++++++++++++++ tests/test_local_backend.py | 21 ++++++++++++++ tests/test_markdown_parser.py | 29 +++++++++++++++++++ tests/test_page_content.py | 40 ++++++++++++++++++++++++++ tests/test_pdf_parser.py | 19 ++++++++++++ tests/test_pipeline.py | 16 +++++++++++ tests/test_review_fixes_2.py | 24 ++++++++++++++++ tests/test_sqlite_storage.py | 54 +++++++++++++++++++++++++++++++++++ tests/test_types.py | 20 +++++++++++++ 18 files changed, 368 insertions(+), 36 deletions(-) create mode 100644 tests/test_types.py diff --git a/pageindex/errors.py b/pageindex/errors.py index f7656ec..aec578a 100644 --- a/pageindex/errors.py +++ b/pageindex/errors.py @@ -8,6 +8,11 @@ class CollectionNotFoundError(PageIndexError): pass +class CollectionAlreadyExistsError(PageIndexError): + """Collection already exists (create_collection, not get_or_create).""" + pass + + class DocumentNotFoundError(PageIndexError): """Document ID not found.""" pass diff --git a/pageindex/index/pipeline.py b/pageindex/index/pipeline.py index f8789a5..f91a587 100644 --- a/pageindex/index/pipeline.py +++ b/pageindex/index/pipeline.py @@ -5,6 +5,13 @@ from ..parser.protocol import ContentNode, ParsedDocument def detect_strategy(nodes: list[ContentNode]) -> str: """Determine which indexing strategy to use based on node data.""" + if not nodes: + # No content at all (e.g. an empty/whitespace-only source file) -> + # level_based's build_tree_from_levels([]) returns an empty structure + # immediately with zero LLM calls. content_based's TOC-detection + # pipeline needs real page content; on an empty page_list it wastes an + # LLM call and then still raises, for no benefit. + return "level_based" if any(n.level is not None for n in nodes): return "level_based" return "content_based" diff --git a/pageindex/page_index.py b/pageindex/page_index.py index 974ce11..0db0c2b 100644 --- a/pageindex/page_index.py +++ b/pageindex/page_index.py @@ -3,6 +3,8 @@ # pageindex/index/page_index.py (the single source of truth). This module # re-exports it so legacy imports (`from pageindex.page_index import ...`, # `from pageindex import page_index`) keep working. +import sys +import types import warnings warnings.warn( @@ -13,3 +15,26 @@ warnings.warn( ) from .index.page_index import * # noqa: F401,F403,E402 + +# pageindex/__init__.py binds the FUNCTION `page_index` as the package +# attribute `pageindex.page_index` (`from .index.page_index import *`). But +# this file is ALSO a real submodule of the same name — the moment anything, +# anywhere in the process, does `import pageindex.page_index` (exactly what +# `from pageindex.page_index import X` triggers), Python's import machinery +# overwrites that package attribute with THIS module object, clobbering the +# function binding. Afterwards `from pageindex import page_index; page_index(x)` +# would raise "TypeError: 'module' object is not callable" — silently, and +# depending entirely on whether this submodule happened to be imported yet. +# +# Fix: make this module itself callable, delegating to the real function, so +# whichever object ends up sitting in the `pageindex.page_index` slot — the +# function or this module — is callable either way. Both `from pageindex.page_index +# import page_index_main` (module attribute access) and +# `from pageindex import page_index; page_index(x)` (call) keep working +# regardless of import order. +class _CallableModule(types.ModuleType): + def __call__(self, *args, **kwargs): + return page_index(*args, **kwargs) + + +sys.modules[__name__].__class__ = _CallableModule diff --git a/pageindex/parser/markdown.py b/pageindex/parser/markdown.py index 7c843e8..04b7e22 100644 --- a/pageindex/parser/markdown.py +++ b/pageindex/parser/markdown.py @@ -12,7 +12,12 @@ class MarkdownParser: path = Path(file_path) model = kwargs.get("model") - with open(path, "r", encoding="utf-8") as f: + # utf-8-sig strips a leading BOM if present (common from Windows + # editors/exporters) and is otherwise identical to plain utf-8. Without + # it, a BOM-prefixed first line fails the header regex below (the BOM + # isn't whitespace, so .strip() doesn't remove it), misclassifying the + # document's first heading as unrecognized preamble text. + with open(path, "r", encoding="utf-8-sig") as f: content = f.read() lines = content.split("\n") @@ -23,7 +28,10 @@ class MarkdownParser: def _extract_headers(self, lines: list[str]) -> list[dict]: header_pattern = r"^(#{1,6})\s+(.+)$" - code_block_pattern = r"^```" + # CommonMark allows both backtick and tilde fences; only recognizing + # backticks let a '#'-prefixed line inside a ~~~-fenced block (e.g. a + # shell comment in a code sample) be misparsed as a real heading. + code_block_pattern = r"^(?:```|~~~)" headers = [] in_code_block = False diff --git a/pageindex/parser/pdf.py b/pageindex/parser/pdf.py index a2739a5..2e02226 100644 --- a/pageindex/parser/pdf.py +++ b/pageindex/parser/pdf.py @@ -77,7 +77,13 @@ class PdfParser: try: pix = pymupdf.Pixmap(image_bytes) - if pix.n > 4: + # n includes the alpha channel, so a plain RGBA pixmap also + # has n==4 — subtract alpha before comparing. Without this, + # a CMYK image with no alpha (n==4, same as RGBA) skips the + # RGB conversion, and pix.save() as .png then raises + # "unsupported colorspace for 'png'", silently dropping the + # image via the bare except below. + if pix.n - pix.alpha >= 4: pix = pymupdf.Pixmap(pymupdf.csRGB, pix) filename = f"p{page_num}_img{img_idx}.png" save_path = images_path / filename diff --git a/pageindex/retrieve.py b/pageindex/retrieve.py index 96d227e..72292eb 100644 --- a/pageindex/retrieve.py +++ b/pageindex/retrieve.py @@ -1,27 +1,24 @@ import json -import PyPDF2 try: - from .index.utils import get_number_of_pages, remove_fields, get_md_page_content + from .index.utils import ( + get_number_of_pages, remove_fields, get_md_page_content, + parse_pages, get_pdf_page_content, + ) except ImportError: - from index.utils import get_number_of_pages, remove_fields, get_md_page_content + from index.utils import ( + get_number_of_pages, remove_fields, get_md_page_content, + parse_pages, get_pdf_page_content, + ) # ── Helpers ────────────────────────────────────────────────────────────────── def _parse_pages(pages: str) -> list[int]: - """Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints.""" - result = [] - for part in pages.split(','): - part = part.strip() - if '-' in part: - start, end = int(part.split('-', 1)[0].strip()), int(part.split('-', 1)[1].strip()) - if start > end: - raise ValueError(f"Invalid range '{part}': start must be <= end") - result.extend(range(start, end + 1)) - else: - result.append(int(part)) - return sorted(set(result)) + """Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints. + Delegates to the canonical implementation so the two never drift again — + this one used to lack the p>=1 filter and the 1000-page DoS cap.""" + return parse_pages(pages) def _count_pages(doc_info: dict) -> int: @@ -34,7 +31,8 @@ def _count_pages(doc_info: dict) -> int: def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: - """Extract text for specific PDF pages (1-indexed). Prefer cached pages, fallback to PDF.""" + """Extract text for specific PDF pages (1-indexed). Prefer cached pages, + else delegate the file-read fallback to the canonical implementation.""" cached_pages = doc_info.get('pages') if cached_pages: page_map = {p['page']: p['content'] for p in cached_pages} @@ -42,15 +40,7 @@ def _get_pdf_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: {'page': p, 'content': page_map[p]} for p in page_nums if p in page_map ] - path = doc_info['path'] - with open(path, 'rb') as f: - pdf_reader = PyPDF2.PdfReader(f) - total = len(pdf_reader.pages) - valid_pages = [p for p in page_nums if 1 <= p <= total] - return [ - {'page': p, 'content': pdf_reader.pages[p - 1].extract_text() or ''} - for p in valid_pages - ] + return get_pdf_page_content(doc_info['path'], page_nums) def _get_md_page_content(doc_info: dict, page_nums: list[int]) -> list[dict]: diff --git a/pageindex/storage/sqlite.py b/pageindex/storage/sqlite.py index 2ed902f..bb0a461 100644 --- a/pageindex/storage/sqlite.py +++ b/pageindex/storage/sqlite.py @@ -1,8 +1,22 @@ import json +import re import sqlite3 import threading from pathlib import Path +from ..errors import CollectionAlreadyExistsError, PageIndexError + +# Mirrors LocalBackend's own collection-name rule. SQLiteStorage enforces this +# itself (not just relying on LocalBackend's pre-check or the schema's CHECK +# constraint below) because it's a public StorageEngine that can be used +# directly, bypassing LocalBackend entirely. +_COLLECTION_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,128}$') + + +def _validate_collection_name(name: str) -> None: + if not _COLLECTION_NAME_RE.match(name): + raise PageIndexError(f"Invalid collection name: {name!r}. Must be 1-128 chars of [a-zA-Z0-9_-].") + class SQLiteStorage: def __init__(self, db_path: str): @@ -48,7 +62,17 @@ class SQLiteStorage: conn.execute("PRAGMA user_version = 1") conn.executescript(""" CREATE TABLE IF NOT EXISTS collections ( - name TEXT PRIMARY KEY CHECK(length(name) <= 128 AND name GLOB '[a-zA-Z0-9_-]*'), + -- GLOB '*' is "any characters", not a regex quantifier over the + -- preceding class — '[a-zA-Z0-9_-]*' alone only constrains the + -- FIRST character. The second GLOB (NOT ... '*[^...]*') checks + -- every remaining character too, so this is real defense-in-depth + -- for direct SQLiteStorage use (bypassing _validate_collection_name + -- above), not just a first-character gate. + name TEXT PRIMARY KEY CHECK( + length(name) BETWEEN 1 AND 128 + AND name GLOB '[a-zA-Z0-9_-]*' + AND name NOT GLOB '*[^a-zA-Z0-9_-]*' + ), created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); CREATE TABLE IF NOT EXISTS documents ( @@ -70,12 +94,17 @@ class SQLiteStorage: conn.commit() def create_collection(self, name: str) -> None: + _validate_collection_name(name) with self._write_lock: conn = self._get_conn() - conn.execute("INSERT INTO collections (name) VALUES (?)", (name,)) + try: + conn.execute("INSERT INTO collections (name) VALUES (?)", (name,)) + except sqlite3.IntegrityError as e: + raise CollectionAlreadyExistsError(f"Collection '{name}' already exists") from e conn.commit() def get_or_create_collection(self, name: str) -> None: + _validate_collection_name(name) with self._write_lock: conn = self._get_conn() conn.execute("INSERT OR IGNORE INTO collections (name) VALUES (?)", (name,)) diff --git a/pageindex/types.py b/pageindex/types.py index c997497..99cf09f 100644 --- a/pageindex/types.py +++ b/pageindex/types.py @@ -14,13 +14,23 @@ class DocumentInfo(TypedDict): doc_type: str -class DocumentDetail(DocumentInfo, total=False): +class _DocumentDetailRequired(DocumentInfo): + """``structure`` is always present — split into its own (default + total=True) base so the total=False below only applies to the genuinely + optional, backend-specific fields below. A single + ``class DocumentDetail(DocumentInfo, total=False): structure: ...`` would + incorrectly mark ``structure`` optional too, since total=False applies to + the whole class body, not just the fields declared after it. + """ + structure: list[dict[str, Any]] + + +class DocumentDetail(_DocumentDetailRequired, total=False): """A document with its tree, as returned by ``get_document()``. ``structure`` is always present; ``file_path`` is local-only and - ``status`` is cloud-only, hence total=False. + ``status`` is cloud-only, hence total=False for those two only. """ - structure: list[dict[str, Any]] file_path: str # local backend only status: str # cloud backend only diff --git a/run_pageindex.py b/run_pageindex.py index d804c6a..c9a0714 100644 --- a/run_pageindex.py +++ b/run_pageindex.py @@ -37,8 +37,8 @@ if __name__ == "__main__": help='Add raw text to nodes (off by default). Bare flag or yes/no') # Markdown specific arguments - parser.add_argument('--if-thinning', type=str, default='no', - help='Whether to apply tree thinning for markdown (markdown only)') + parser.add_argument('--if-thinning', nargs='?', const=True, type=_cli_bool, default=None, + help='Apply tree thinning (off by default, markdown only). Bare flag or yes/no') parser.add_argument('--thinning-threshold', type=int, default=5000, help='Minimum token threshold for thinning (markdown only)') parser.add_argument('--summary-token-threshold', type=int, default=200, @@ -101,7 +101,7 @@ if __name__ == "__main__": toc_with_page_number = asyncio.run(md_to_tree( md_path=args.md_path, - if_thinning=args.if_thinning.lower() == 'yes', + if_thinning=bool(args.if_thinning), min_token_threshold=args.thinning_threshold, if_add_node_summary=opt.if_add_node_summary, summary_token_threshold=args.summary_token_threshold, diff --git a/tests/test_legacy_shims.py b/tests/test_legacy_shims.py index c94fbb1..d24e2f3 100644 --- a/tests/test_legacy_shims.py +++ b/tests/test_legacy_shims.py @@ -3,10 +3,15 @@ now deprecation shims over the canonical pageindex.index.* modules. These tests pin the compatibility contract.""" import asyncio import importlib +import subprocess +import sys import warnings +from pathlib import Path import pytest +_REPO_ROOT = Path(__file__).resolve().parent.parent + def test_plain_import_pageindex_does_not_warn(): # `import pageindex` must not route through the deprecation shims. @@ -116,3 +121,27 @@ def test_md_to_tree_coerces_legacy_yes_no_strings(tmp_path): assert not _has_summary(result["structure"]) assert all("node_id" in n for n in result["structure"]) + + +def test_page_index_stays_callable_after_the_submodule_is_imported(): + """pageindex/__init__.py binds the FUNCTION `page_index` as the package + attribute, but pageindex/page_index.py is ALSO a real submodule of the + same name — importing that submodule anywhere clobbers the package + attribute with the module object (Python's import machinery does this + unconditionally). Must run in a fresh subprocess: the effect depends on + import order, so it can't be reliably observed against an + already-imported pageindex in this test process.""" + script = ( + "import warnings; warnings.simplefilter('ignore')\n" + "import pageindex.page_index\n" # the clobbering import + "from pageindex import page_index\n" + "assert callable(page_index), f'page_index is not callable: {type(page_index)}'\n" + "from pageindex.page_index import page_index_main\n" # old multi-symbol import still works + "assert callable(page_index_main)\n" + "print('OK')\n" + ) + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, cwd=str(_REPO_ROOT), + ) + assert result.returncode == 0, result.stderr + assert "OK" in result.stdout diff --git a/tests/test_local_backend.py b/tests/test_local_backend.py index 0b2f1f3..72b4129 100644 --- a/tests/test_local_backend.py +++ b/tests/test_local_backend.py @@ -35,6 +35,27 @@ def test_unsupported_file_type_raises(backend, tmp_path): backend.add_document("papers", str(bad_file)) +def test_add_document_on_empty_markdown_file_does_not_crash(tmp_path): + """An empty/whitespace-only .md file used to route into the PDF-oriented + TOC-detection pipeline (no node ever has 'level' set), wasting an LLM call + and then raising IndexingError. Must complete instantly with zero LLM + calls when summary/description are off.""" + from pageindex.config import IndexConfig + + storage = SQLiteStorage(str(tmp_path / "test.db")) + backend = LocalBackend( + storage=storage, files_dir=str(tmp_path / "files"), model="gpt-4o", + index_config=IndexConfig(if_add_node_summary=False, if_add_doc_description=False), + ) + backend.get_or_create_collection("papers") + empty_md = tmp_path / "empty.md" + empty_md.write_text(" \n\n \n") + + doc_id = backend.add_document("papers", str(empty_md)) # must not raise + + assert backend.get_document_structure("papers", doc_id) == [] + + def test_register_custom_parser(backend): from pageindex.parser.protocol import ParsedDocument, ContentNode diff --git a/tests/test_markdown_parser.py b/tests/test_markdown_parser.py index 6d337c1..bffda98 100644 --- a/tests/test_markdown_parser.py +++ b/tests/test_markdown_parser.py @@ -71,3 +71,32 @@ def test_headerless_file_yields_single_node(tmp_path): assert len(result.nodes) == 1 assert result.nodes[0].title == "plain" assert "No headings at all" in result.nodes[0].content + + +def test_utf8_bom_does_not_break_the_first_header(tmp_path): + """A leading BOM (common from Windows editors/exporters) isn't + whitespace, so .strip() doesn't remove it — without utf-8-sig decoding, + the header regex fails to match the BOM-prefixed first line, and it gets + misclassified as unrecognized preamble text instead of a real heading.""" + md = tmp_path / "bom.md" + md.write_bytes(b"\xef\xbb\xbf# First Header\nbody text\n") + result = MarkdownParser().parse(str(md)) + assert len(result.nodes) == 1 + assert result.nodes[0].title == "First Header" + assert result.nodes[0].level == 1 + + +def test_tilde_fenced_code_blocks_are_recognized(tmp_path): + """CommonMark allows both backtick and tilde code fences. Only + recognizing backticks let a '#'-prefixed line inside a ~~~-fenced block + (e.g. a shell comment in a sample) be misparsed as a real heading.""" + md = tmp_path / "tilde.md" + md.write_text( + "# Real Header\nintro\n" + "~~~\n# not a real header, just a comment\n~~~\n" + "## Real Sub\nmore\n" + ) + result = MarkdownParser().parse(str(md)) + titles = [n.title for n in result.nodes] + assert titles == ["Real Header", "Real Sub"] + assert "not a real header" not in " ".join(n.title for n in result.nodes) diff --git a/tests/test_page_content.py b/tests/test_page_content.py index e810651..c610180 100644 --- a/tests/test_page_content.py +++ b/tests/test_page_content.py @@ -34,3 +34,43 @@ def test_retrieve_md_page_content_returns_only_requested_lines(): out = _get_md_page_content({"structure": _md_structure()}, [5, 100]) assert [r["page"] for r in out] == [5, 100] + + +def test_retrieve_parse_pages_delegates_to_canonical_and_enforces_dos_cap(): + """retrieve._parse_pages used to be an independent copy that lacked the + canonical parse_pages' p>=1 filter and 1000-page cap — a caller of the + legacy pageindex.get_page_content could bypass the DoS guard the SDK path + enforces. Now it's a one-line delegate, so they can't drift again.""" + from pageindex.retrieve import _parse_pages + from pageindex.index.utils import parse_pages + import pytest + + assert _parse_pages("5-7") == parse_pages("5-7") == [5, 6, 7] + with pytest.raises(ValueError, match="too large"): + _parse_pages("1-99999999") + + +def test_retrieve_get_pdf_page_content_falls_back_to_canonical(tmp_path, monkeypatch): + """When no cached 'pages' are present, the file-read fallback must + delegate to the canonical get_pdf_page_content instead of re-implementing + PDF text extraction inline (a second, independently-maintained copy).""" + from pageindex.retrieve import _get_pdf_page_content + import pageindex.retrieve as retrieve_mod + + calls = [] + monkeypatch.setattr( + retrieve_mod, "get_pdf_page_content", + lambda path, page_nums: calls.append((path, page_nums)) or [{"page": 1, "content": "x"}], + ) + result = _get_pdf_page_content({"path": "/fake/doc.pdf"}, [1]) + assert calls == [("/fake/doc.pdf", [1])] + assert result == [{"page": 1, "content": "x"}] + + +def test_retrieve_get_pdf_page_content_prefers_cache_over_file(): + from pageindex.retrieve import _get_pdf_page_content + + doc_info = {"path": "/should/not/be/opened.pdf", + "pages": [{"page": 1, "content": "cached one"}, {"page": 2, "content": "cached two"}]} + result = _get_pdf_page_content(doc_info, [2]) + assert result == [{"page": 2, "content": "cached two"}] diff --git a/tests/test_pdf_parser.py b/tests/test_pdf_parser.py index 0d6e051..cafe556 100644 --- a/tests/test_pdf_parser.py +++ b/tests/test_pdf_parser.py @@ -1,3 +1,4 @@ +import pymupdf import pytest from pathlib import Path from pageindex.parser.pdf import PdfParser @@ -29,6 +30,24 @@ def test_parse_nodes_are_flat_without_level(): assert node.level is None +def test_cmyk_pixmap_without_alpha_is_saveable_as_png(tmp_path): + """A CMYK image with no alpha has n==4 -- same as RGBA -- so `pix.n > 4` + wrongly skips the RGB conversion PNG needs, and pix.save() raises + 'unsupported colorspace for png', silently dropping the image via the + extractor's bare except. The fix (`pix.n - pix.alpha >= 4`) must convert + CMYK (4-0=4) while leaving RGBA (4-1=3) untouched.""" + cmyk = pymupdf.Pixmap(pymupdf.csCMYK, pymupdf.Rect(0, 0, 10, 10)) + assert cmyk.n == 4 and cmyk.alpha == 0 + assert cmyk.n - cmyk.alpha >= 4 # the fixed condition: must convert + converted = pymupdf.Pixmap(pymupdf.csRGB, cmyk) + converted.save(str(tmp_path / "cmyk.png")) # must not raise + + rgba = pymupdf.Pixmap(pymupdf.Pixmap(pymupdf.csRGB, pymupdf.Rect(0, 0, 10, 10)), 1) + assert rgba.n == 4 and rgba.alpha == 1 + assert not (rgba.n - rgba.alpha >= 4) # unchanged: RGBA needs no conversion + rgba.save(str(tmp_path / "rgba.png")) # already saveable as-is + + def test_image_paths_are_absolute(tmp_path): """Image references must be absolute so they resolve regardless of cwd (cwd-relative paths broke after the query ran from another directory).""" diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 6b0cf12..98a66ba 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -25,6 +25,22 @@ def test_detect_strategy_without_level(): assert detect_strategy(nodes) == "content_based" +def test_detect_strategy_empty_nodes_is_level_based(): + """An empty node list (e.g. an empty/whitespace-only source file) must + route to level_based, whose build_tree_from_levels([]) returns an empty + structure with zero LLM calls — not content_based, whose TOC-detection + pipeline needs real page content and wastes an LLM call before failing.""" + assert detect_strategy([]) == "level_based" + + +def test_build_index_on_empty_document_makes_no_llm_calls(): + from pageindex.config import IndexConfig + parsed = ParsedDocument(doc_name="empty", nodes=[]) + opt = IndexConfig(if_add_node_summary=False, if_add_doc_description=False) + result = build_index(parsed, opt=opt) + assert result == {"doc_name": "empty", "structure": []} + + def test_build_tree_from_levels(): nodes = [ ContentNode(content="ch1 text", tokens=10, title="Chapter 1", index=1, level=1), diff --git a/tests/test_review_fixes_2.py b/tests/test_review_fixes_2.py index 5bfb903..e1dbbcf 100644 --- a/tests/test_review_fixes_2.py +++ b/tests/test_review_fixes_2.py @@ -2,6 +2,7 @@ 2d46d68..8f536cb): the Markdown text-stripping fix's fallout, plus the other directly-fixable findings from that pass.""" import asyncio +from unittest.mock import AsyncMock import pytest @@ -122,6 +123,20 @@ def test_cloud_delete_collection_still_clears_real_folder_id(monkeypatch): # ── #7: remove_structure_text is skipped when text was never added ─────────── +def _mock_content_based_pipeline(monkeypatch, structure): + """content_based's real path (_content_based_pipeline) drives real LLM + calls (TOC detection etc.) regardless of if_add_node_summary — a prior + version of these two tests didn't mock this out, fell through to it, and + made real network calls (with a dummy key: 10 retries before failing; + with a real key: real billable requests) on every run.""" + from pageindex.index import pipeline + + async def fake(page_list, opt): + return structure + + monkeypatch.setattr(pipeline, "_content_based_pipeline", fake) + + def test_build_index_skips_text_strip_when_no_text_was_added(monkeypatch): from pageindex.index import pipeline from pageindex.parser.protocol import ContentNode, ParsedDocument @@ -131,6 +146,7 @@ def test_build_index_skips_text_strip_when_no_text_was_added(monkeypatch): # ...` inside the function body), so patch it on the utils module itself. import pageindex.index.utils as utils_mod monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s) + _mock_content_based_pipeline(monkeypatch, [{"title": "T", "start_index": 1, "end_index": 1}]) nodes = [ContentNode(content="page one text", tokens=5, index=1)] parsed = ParsedDocument(doc_name="d", nodes=nodes) @@ -147,6 +163,14 @@ def test_build_index_still_strips_text_when_summary_added_it(monkeypatch): calls = [] import pageindex.index.utils as utils_mod monkeypatch.setattr(utils_mod, "remove_structure_text", lambda s: calls.append(s) or s) + _mock_content_based_pipeline(monkeypatch, [{"title": "T", "start_index": 1, "end_index": 1}]) + # Summary generation itself would otherwise make a real LLM call. + monkeypatch.setattr( + utils_mod, "generate_summaries_for_structure", + AsyncMock(side_effect=lambda structure, model=None: [ + n.__setitem__("summary", "fake") for n in structure + ]), + ) nodes = [ContentNode(content="page one text", tokens=5, index=1)] parsed = ParsedDocument(doc_name="d", nodes=nodes) diff --git a/tests/test_sqlite_storage.py b/tests/test_sqlite_storage.py index c21da94..3fa2a43 100644 --- a/tests/test_sqlite_storage.py +++ b/tests/test_sqlite_storage.py @@ -19,6 +19,60 @@ def test_delete_collection(storage): storage.delete_collection("papers") assert "papers" not in storage.list_collections() + +def test_create_duplicate_collection_raises_pageindex_error(storage): + """A raw sqlite3.IntegrityError leaking out breaks `except PageIndexError` + catch-alls; must be translated to a proper SDK exception.""" + from pageindex.errors import CollectionAlreadyExistsError, PageIndexError + storage.create_collection("papers") + with pytest.raises(CollectionAlreadyExistsError): + storage.create_collection("papers") + # also catchable via the SDK's generic base class + storage.create_collection("other") + with pytest.raises(PageIndexError): + storage.create_collection("other") + + +@pytest.mark.parametrize("bad_name", [ + "a/../../etc/passwd", "/etc/passwd", "a$(whoami)", ".hidden", + "a b", "válid", "", "a" * 129, +]) +def test_create_collection_rejects_invalid_names_at_the_python_layer(storage, bad_name): + """SQLiteStorage must validate collection names itself — it's a public + StorageEngine that can be used directly, bypassing LocalBackend's own + regex check entirely.""" + from pageindex.errors import PageIndexError + with pytest.raises(PageIndexError): + storage.create_collection(bad_name) + + +def test_sql_check_constraint_also_rejects_invalid_names_directly(storage): + """Defense-in-depth: even bypassing SQLiteStorage's own Python validation + and inserting via raw SQL, the schema's CHECK constraint must reject a + name that isn't ENTIRELY [a-zA-Z0-9_-] — not just its first character + (GLOB '*' is a wildcard, not a regex quantifier over the preceding class, + so 'name GLOB [a-zA-Z0-9_-]*' alone only constrains the first character).""" + import sqlite3 + conn = storage._get_conn() + with pytest.raises(sqlite3.IntegrityError): + conn.execute("INSERT INTO collections (name) VALUES (?)", ("a/../../etc/passwd",)) + + +def test_malicious_collection_name_rejected_through_local_backend_too(tmp_path): + """End-to-end via the normal LocalBackend entry point: a path-traversal- + shaped collection name must never reach add_document's + files_dir / collection path construction. Three independent layers now + reject it (LocalBackend's own regex, SQLiteStorage's regex, and the SQL + CHECK constraint) — this pins the outermost one.""" + from pageindex.backend.local import LocalBackend + from pageindex.errors import PageIndexError + + storage = SQLiteStorage(str(tmp_path / "t.db")) + backend = LocalBackend(storage=storage, files_dir=str(tmp_path / "files"), model="gpt-4o") + with pytest.raises(PageIndexError): + backend.create_collection("a/../../escape_me") + assert not (tmp_path / "escape_me").exists() + def test_save_and_get_document(storage): storage.create_collection("papers") doc = { diff --git a/tests/test_types.py b/tests/test_types.py new file mode 100644 index 0000000..093565e --- /dev/null +++ b/tests/test_types.py @@ -0,0 +1,20 @@ +from pageindex.types import DocumentDetail, DocumentInfo, PageContent + + +def test_document_detail_structure_field_is_required(): + """structure is always populated by both LocalBackend.get_document and + CloudBackend.get_document — must be a required key, not optional, or + type checkers/tooling built on this TypedDict wrongly treat a + DocumentDetail missing 'structure' as valid.""" + assert "structure" in DocumentDetail.__required_keys__ + assert "structure" not in DocumentDetail.__optional_keys__ + + +def test_document_detail_backend_specific_fields_stay_optional(): + assert "file_path" in DocumentDetail.__optional_keys__ + assert "status" in DocumentDetail.__optional_keys__ + + +def test_document_detail_inherits_document_info_as_required(): + for key in ("doc_id", "doc_name", "doc_description", "doc_type"): + assert key in DocumentDetail.__required_keys__