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
This commit is contained in:
mountain 2026-07-09 11:58:59 +08:00
parent b9d021916f
commit 4e6a13576d
18 changed files with 368 additions and 36 deletions

View file

@ -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

View file

@ -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"

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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]:

View file

@ -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,))

View file

@ -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

View file

@ -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,

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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"}]

View file

@ -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)."""

View file

@ -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),

View file

@ -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)

View file

@ -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 = {

20
tests/test_types.py Normal file
View file

@ -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__