126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
"""Walker — iterate .state/comments/, look up inline body, parallelize."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import fitz
|
|
|
|
from rex.comments.walker import walk_and_extract
|
|
|
|
|
|
def _make_pdf(path: Path, text: str = "x" * 200) -> None:
|
|
doc = fitz.open()
|
|
doc.new_page().insert_text((50, 72), text)
|
|
doc.save(str(path))
|
|
doc.close()
|
|
|
|
|
|
def _setup_two_comments(root: Path) -> tuple[Path, Path]:
|
|
a = root / "CMS-2024-0001" / "CMS-2024-0001-0001"
|
|
b = root / "CMS-2024-0001" / "CMS-2024-0001-0002"
|
|
a.mkdir(parents=True)
|
|
b.mkdir(parents=True)
|
|
_make_pdf(a / "attachment_1.pdf", "Comment A body content. " * 10)
|
|
_make_pdf(b / "attachment_1.pdf", "Comment B body content. " * 10)
|
|
return a, b
|
|
|
|
|
|
def _stub_inline_body(_cid: str) -> str:
|
|
return "inline body for stub"
|
|
|
|
|
|
def test_walk_and_extract_writes_all_combined(tmp_path: Path):
|
|
a, b = _setup_two_comments(tmp_path)
|
|
|
|
stats = walk_and_extract(tmp_path, inline_body_lookup=_stub_inline_body, workers=2)
|
|
|
|
assert (a / "combined.md").is_file()
|
|
assert (b / "combined.md").is_file()
|
|
assert stats["written"] == 2
|
|
assert stats["skipped"] == 0
|
|
|
|
|
|
def test_walk_and_extract_skips_existing(tmp_path: Path):
|
|
a, _b = _setup_two_comments(tmp_path)
|
|
(a / "combined.md").write_text("---\ncomment_id: x\n---\n\npre-existing\n")
|
|
os.utime(a / "attachment_1.pdf", ns=(1_000, 1_000))
|
|
|
|
stats = walk_and_extract(tmp_path, inline_body_lookup=_stub_inline_body, workers=1)
|
|
|
|
assert stats["written"] == 1 # only the second one
|
|
assert stats["skipped"] == 1
|
|
assert "pre-existing" in (a / "combined.md").read_text() # untouched
|
|
|
|
|
|
def test_walk_and_extract_filters_by_docket(tmp_path: Path):
|
|
_setup_two_comments(tmp_path)
|
|
other = tmp_path / "CMS-2099-9999" / "CMS-2099-9999-0001"
|
|
other.mkdir(parents=True)
|
|
_make_pdf(other / "attachment_1.pdf", "other docket content. " * 10)
|
|
|
|
stats = walk_and_extract(
|
|
tmp_path,
|
|
inline_body_lookup=_stub_inline_body,
|
|
docket="CMS-2099-9999",
|
|
workers=1,
|
|
)
|
|
|
|
assert stats["written"] == 1
|
|
assert (other / "combined.md").is_file()
|
|
|
|
|
|
def test_walk_and_extract_respects_limit(tmp_path: Path):
|
|
_setup_two_comments(tmp_path)
|
|
stats = walk_and_extract(
|
|
tmp_path, inline_body_lookup=_stub_inline_body, limit=1, workers=1
|
|
)
|
|
assert stats["written"] == 1
|
|
|
|
|
|
def test_on_extracted_called_for_written_dirs(tmp_path: Path):
|
|
a, b = _setup_two_comments(tmp_path)
|
|
seen: list[tuple[str, Path]] = []
|
|
|
|
walk_and_extract(
|
|
tmp_path,
|
|
inline_body_lookup=_stub_inline_body,
|
|
workers=2,
|
|
on_extracted=lambda cid, d: seen.append((cid, d)),
|
|
)
|
|
|
|
assert sorted(cid for cid, _ in seen) == sorted([a.name, b.name])
|
|
# callback gets the comment dir; sibling + combined.md are inside it
|
|
for _cid, d in seen:
|
|
assert (d / "combined.md").is_file()
|
|
assert (d / "attachment_1.pdf.md").is_file()
|
|
|
|
|
|
def test_on_extracted_called_for_skipped_dirs(tmp_path: Path):
|
|
"""With reattach=True, pre-existing combined.md still gets the callback,
|
|
and any missing sibling MDs are derived from it before the callback
|
|
fires — so previously-extracted dirs end up with the same set of MDs
|
|
as freshly-written ones."""
|
|
a, b = _setup_two_comments(tmp_path)
|
|
# Pre-write a stale combined.md with one section so derive_siblings
|
|
# has something to split.
|
|
(a / "combined.md").write_text(
|
|
"---\ncomment_id: x\ndocket_id: y\n---\n\n## attachment_1.pdf\n\nstale body\n"
|
|
)
|
|
|
|
seen: list[str] = []
|
|
stats = walk_and_extract(
|
|
tmp_path,
|
|
inline_body_lookup=_stub_inline_body,
|
|
workers=1,
|
|
reattach=True,
|
|
on_extracted=lambda cid, _d: seen.append(cid),
|
|
)
|
|
|
|
assert stats["written"] == 1 # b
|
|
assert stats["skipped"] == 1 # a
|
|
assert sorted(seen) == sorted([a.name, b.name])
|
|
# a's sibling was backfilled from its (stale) combined.md
|
|
assert (a / "attachment_1.pdf.md").is_file()
|