tests/test_ocr_release.py

"""Tiny real SQLite/text candidates; never open/rebuild production databases."""
import copy
import gzip
import hashlib
import json
from pathlib import Path
import tempfile
import unittest
from unittest.mock import patch

import app
import build_fulltext_db as builder
import ocr_release as release
from fastapi.testclient import TestClient


class OcrReleaseTests(unittest.TestCase):
    def setUp(self):
        temporary = tempfile.TemporaryDirectory()
        self.addCleanup(temporary.cleanup)
        self.root = Path(temporary.name)
        self.bucket = self.root / "bucket"
        self.bucket.mkdir()
        self.base = self.root / "base.json.gz"
        base_record = {"doc_id": "CCRD:1", "source": "CCRD", "source_name": "CCRD",
                       "storage_root": "CCRD", "storage_rel_path": "old.txt", "display_name": "旧文件",
                       "display_dirs": [], "display_rel_path": "旧文件.txt", "size": 7, "extension": "txt"}
        self.base.write_bytes(gzip.compress(json.dumps({"records": [base_record], "sources": [
            {"slug": "CCRD", "name": "CCRD", "count": 1}]}).encode(), mtime=0))
        self.base_sha = hashlib.sha256(self.base.read_bytes()).hexdigest()
        self.content = "手机文化\n\n\n\n末页手。机 ABC"
        texts = ["手机文化", "", "末页手。机 ABC"]
        body = self.content.encode()
        txt_sha = release.digest(body)
        text_ref = {"bucket": release.BUCKET, "path": f"text/objects/{txt_sha[:2]}/{txt_sha}.txt",
                    "sha256": txt_sha, "bytes": len(body), "role": "runtime"}
        self.write(text_ref["path"], body)
        original = "原目录/book.pdf"
        identity = {"recipe": "ocr-plain-utf8-pages-v1", "origin_repo": "VoiceOfML/Test", "origin_rel_path": original,
                    "origin_revision": "b" * 40, "origin_sha256": "a" * 64,
                    "source_sha256": "a" * 64, "doc_id": "OCR:1",
                    "text_generation": "c" * 64, "txt_sha256": txt_sha}
        generation = release.digest(release.canonical(identity).rstrip(b"\n"))
        book_id = release.digest(("VoiceOfML/Test\0" + original).encode())
        pages, cursor = [], 0
        for number, text in enumerate(texts, 1):
            pages.append({"page": number, "start": cursor, "end": cursor + len(text),
                          "text_sha256": release.digest(text.encode()), "text_generation": str(number) * 64,
                          "quality": "partially-reviewed" if number == 1 else "unreviewed"})
            cursor += len(text) + 2
        self.export = {"version": 1, "kind": "ccrd-ocr-text-export", "complete": True, **identity,
                       "generation": generation, "book_id": book_id, "title": "book", "page_count": 3,
                       "source_sha256_role": "reader-primary",
                       "readable_rel_path": "VoiceOfML/Test/原目录/book.txt", "accepted_pages": 1,
                       "quality": "partially-reviewed", "offset_unit": "unicode-codepoint", "page_separator": "\n\n",
                       "pages": pages, "text": text_ref,
                       "retained_reader_resources": [{"bucket": "vomebook/pdf-pages-v2", "path": "objects/aa/reading-manifest.json",
                                                      "sha256": "a" * 64, "bytes": 1, "role": "provenance"}]}
        self.export_path = f"text/exports/{book_id}/{generation}/manifest.json"
        self.export_ref = self.ref(self.export_path, release.canonical(self.export))
        core = {"version": 1, "kind": "ccrd-ocr-search-release", "base_metadata_sha256": self.base_sha,
                "builder_sha256": release.digest(Path(builder.__file__).read_bytes()), "tokenizer_version": release.TOKENIZER,
                "exports": [self.export_ref]}
        generation = release.digest(release.canonical(core))
        records = [release.export_record(self.bucket, self.export_ref)[0]]
        db_relative = f"indexes/ocr/{generation}/OCR.sqlite3"
        directory = self.bucket / Path(db_relative).parent
        directory.mkdir(parents=True)
        with patch.object(builder, "BASE_DIR", self.bucket), patch.object(builder, "FULLTEXT_DIR", directory):
            builder.build_source_db("OCR", records)
        self.manifest = {**core, "generation": generation, "documents": 1,
                         "database": self.ref(db_relative, (self.bucket / db_relative).read_bytes())}
        self.manifest_path = f"releases/ocr-{generation}/manifest.json"
        raw = release.canonical(self.manifest)
        self.write(self.manifest_path, raw)
        self.manifest_sha = release.digest(raw)

    def write(self, path, body):
        target = self.bucket / path
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(body)

    def ref(self, path, body):
        self.write(path, body)
        return {"bucket": release.BUCKET, "path": path, "sha256": release.digest(body), "bytes": len(body), "role": "runtime"}

    def load(self):
        return release.load_release(self.bucket, self.manifest_path, self.manifest_sha, self.base_sha)

    def test_whole_release_is_verified_and_origin_hierarchy_is_preserved(self):
        candidate = self.load()
        self.assertEqual(candidate["records"][0]["display_rel_path"], "VoiceOfML/Test/原目录/book.txt")
        self.assertEqual(candidate["text_paths"]["OCR:1"].read_text(), self.content)
        self.assertEqual(candidate["tree"][0]["count"], 1)
        self.assertEqual(candidate["browser"]["VoiceOfML/Test/原目录"]["files"][0]["doc_id"], "OCR:1")

    def test_corrupt_text_or_database_blocks_the_release(self):
        for ref in (self.export["text"], self.manifest["database"]):
            target = self.bucket / ref["path"]
            original = target.read_bytes()
            try:
                target.write_bytes(original + b"damage")
                with self.assertRaisesRegex(ValueError, "checksum"):
                    self.load()
            finally:
                target.write_bytes(original)

    def test_wrong_manifest_or_base_identity_cannot_activate(self):
        for manifest_sha, base_sha in (("d" * 64, self.base_sha), (self.manifest_sha, "e" * 64)):
            with self.subTest(manifest_sha=manifest_sha), self.assertRaises(ValueError):
                release.load_release(self.bucket, self.manifest_path, manifest_sha, base_sha)

    def test_export_page_offsets_and_missing_provenance_are_rejected(self):
        for change in ({"pages": []}, {"retained_reader_resources": []}, {"readable_rel_path": "wrong/book.txt"},
                       {"accepted_pages": 3}, {"origin_repo": "../bad"}):
            value = {**copy.deepcopy(self.export), **change}
            ref = self.ref(self.export_path, release.canonical(value))
            with self.subTest(change=list(change)), self.assertRaises(ValueError):
                release.export_record(self.bucket, ref)

    def test_pin_copies_verified_data_and_is_independent_of_the_mount(self):
        pinned = release.pin_release(self.bucket, self.manifest_path, self.manifest_sha, self.base_sha, self.root / "cache")
        self.assertNotIn(self.bucket, pinned["database"].parents)
        self.assertEqual(pinned["text_paths"]["OCR:1"].read_text(), self.content)
        self.write(self.export["text"]["path"], b"changed remote mount")
        self.assertEqual(pinned["text_paths"]["OCR:1"].read_text(), self.content)

    def test_complete_body_admission_binds_each_page_even_when_summary_says_passed(self):
        base=copy.deepcopy(self.export)
        admission={'policy':'complete-body-quality-v1','source_sha256':base['source_sha256'],
                   'text_generation':base['text_generation'],
                   'quality_check':{'status':'passed','issues':[],'checked_pages':3,'page_count':3},
                   'pages':[{'page':n,'quality_check':{'page':n,'status':'passed','issues':[]}} for n in [1,2,3]]}
        def write(value):
            export={**base,'body_admission':value,'body_admission_sha256':release.digest(release.canonical(value).rstrip(b'\n'))}
            identity={k:export[k] for k in ('recipe','origin_repo','origin_rel_path','origin_revision','origin_sha256',
                       'source_sha256','doc_id','text_generation','txt_sha256','body_admission_sha256')}
            export['generation']=release.digest(release.canonical(identity).rstrip(b'\n'))
            path=f'text/exports/{export["book_id"]}/{export["generation"]}/manifest.json'
            return self.ref(path,release.canonical(export))
        record,export=release.export_record(self.bucket,write(admission))
        self.assertEqual(record['doc_id'],'OCR:1')
        self.assertEqual(export['body_admission'],admission)
        for change in ['status','issue','missing','duplicate']:
            broken=copy.deepcopy(admission)
            if change=='status':broken['pages'][1]['quality_check']['status']='review'
            elif change=='issue':broken['pages'][1]['quality_check']['issues']=[{'code':'uncertain'}]
            elif change=='missing':broken['pages'].pop()
            else:broken['pages'][1]['page']=1
            with self.subTest(change=change),self.assertRaisesRegex(ValueError,'complete-body quality'):
                release.export_record(self.bucket,write(broken))

    def test_small_index_uses_existing_in_memory_search_contract(self):
        candidate = self.load()
        databases = app.FulltextDatabases(self.root / "empty-db", in_memory=True, extra_paths=[candidate["database"]])
        try:
            self.assertEqual(databases.search_source("OCR", "手机", exact=False), {1})
            self.assertEqual(databases.search_source("OCR", "手。机", exact=True), {1})
            self.assertEqual(databases.search_source("OCR", "文化 ABC", exact=False), {1})
            self.assertEqual(databases.search_source("OCR", "手机 ABC", exact=True), set())
            self.assertEqual(databases.connections["OCR"].execute("PRAGMA database_list").fetchone()[2], "")
        finally:
            databases.close()

    def test_public_layout_splits_repositories_and_rebases_paths(self):
        records = [
            {"doc_id": "OCR:1", "source": "OCR", "source_name": "VoiceOfML", "display_name": "A",
             "display_dirs": ["VoiceOfML", "Zeta", "d1"],
             "display_rel_path": "VoiceOfML/Zeta/d1/A.txt", "extension": "txt", "size": 1},
            {"doc_id": "OCR:2", "source": "OCR", "source_name": "VoiceOfML", "display_name": "B",
             "display_dirs": ["VoiceOfML", "Alpha"],
             "display_rel_path": "VoiceOfML/Alpha/B.txt", "extension": "txt", "size": 1},
        ]
        extra = {
            "records": records,
            "tree": [{
                "name": "VoiceOfML", "path": "VoiceOfML", "count": 2, "hasDirectFiles": False,
                "hasChildren": True, "showSelfToggle": False, "children": [
                    {"name": "Alpha", "path": "VoiceOfML/Alpha", "count": 1, "hasDirectFiles": True,
                     "hasChildren": False, "showSelfToggle": False, "children": []},
                    {"name": "Zeta", "path": "VoiceOfML/Zeta", "count": 1, "hasDirectFiles": False,
                     "hasChildren": True, "showSelfToggle": False, "children": [
                         {"name": "d1", "path": "VoiceOfML/Zeta/d1", "count": 1, "hasDirectFiles": True,
                          "hasChildren": False, "showSelfToggle": False, "children": []}]},
                ]}],
            "browser": {
                "": {"folders": [{"name": "VoiceOfML", "path": "VoiceOfML", "count": 2}], "files": []},
                "VoiceOfML": {"folders": [
                    {"name": "Alpha", "path": "VoiceOfML/Alpha", "count": 1},
                    {"name": "Zeta", "path": "VoiceOfML/Zeta", "count": 1}], "files": []},
                "VoiceOfML/Alpha": {"folders": [], "files": [
                    {"name": "B", "ext": "txt", "size": 1, "doc_id": "OCR:2"}]},
                "VoiceOfML/Zeta": {"folders": [
                    {"name": "d1", "path": "VoiceOfML/Zeta/d1", "count": 1}], "files": []},
                "VoiceOfML/Zeta/d1": {"folders": [], "files": [
                    {"name": "A", "ext": "txt", "size": 1, "doc_id": "OCR:1"}]},
            },
        }
        layout = app.build_ocr_public_layout(extra, {"CCRD"})
        self.assertEqual([source["slug"] for source in layout["sources"]], ["Alpha", "Zeta"])
        self.assertEqual(layout["physical_map"], {"Alpha": "OCR", "Zeta": "OCR"})
        self.assertEqual(layout["trees"]["Zeta"][0]["path"], "")
        self.assertEqual(layout["trees"]["Zeta"][0]["children"][0]["path"], "d1")
        self.assertEqual(layout["browsers"]["Zeta"][""]["folders"][0]["path"], "d1")
        self.assertEqual(records[0]["public_source"], "Zeta")
        self.assertEqual(records[0]["public_dirs"], ["d1"])
        self.assertEqual(records[0]["public_rel_path"], "Zeta/d1/A.txt")
        self.assertEqual(records[1]["public_source"], "Alpha")
        self.assertEqual(records[1]["public_dirs"], [])
        self.assertEqual(records[1]["public_rel_path"], "Alpha/B.txt")

    def test_public_layout_filters_search_by_repository(self):
        records = [
            {"doc_id": "OCR:1", "source": "OCR", "source_name": "VoiceOfML", "display_name": "A",
             "display_dirs": ["VoiceOfML", "Zeta", "d1"],
             "display_rel_path": "VoiceOfML/Zeta/d1/A.txt", "extension": "txt", "size": 1},
            {"doc_id": "OCR:2", "source": "OCR", "source_name": "VoiceOfML", "display_name": "B",
             "display_dirs": ["VoiceOfML", "Alpha"],
             "display_rel_path": "VoiceOfML/Alpha/B.txt", "extension": "txt", "size": 1},
        ]
        extra = {
            "records": records,
            "tree": [{
                "name": "VoiceOfML", "path": "VoiceOfML", "count": 2, "hasDirectFiles": False,
                "hasChildren": True, "showSelfToggle": False, "children": [
                    {"name": "Alpha", "path": "VoiceOfML/Alpha", "count": 1, "hasDirectFiles": True,
                     "hasChildren": False, "showSelfToggle": False, "children": []},
                    {"name": "Zeta", "path": "VoiceOfML/Zeta", "count": 1, "hasDirectFiles": False,
                     "hasChildren": True, "showSelfToggle": False, "children": [
                         {"name": "d1", "path": "VoiceOfML/Zeta/d1", "count": 1, "hasDirectFiles": True,
                          "hasChildren": False, "showSelfToggle": False, "children": []}]},
                ]}],
            "browser": {
                "VoiceOfML/Alpha": {"folders": [], "files": [
                    {"name": "B", "ext": "txt", "size": 1, "doc_id": "OCR:2"}]},
                "VoiceOfML/Zeta": {"folders": [
                    {"name": "d1", "path": "VoiceOfML/Zeta/d1", "count": 1}], "files": []},
                "VoiceOfML/Zeta/d1": {"folders": [], "files": [
                    {"name": "A", "ext": "txt", "size": 1, "doc_id": "OCR:1"}]},
            },
        }
        layout = app.build_ocr_public_layout(extra, set())
        saved = {key: getattr(app, key) for key in (
            "records", "source_counts", "source_physical_map", "physical_source_counts",
            "ocr_public_browser", "folder_tree_data", "folder_browser_data", "fulltext_databases",
            "source_records_map", "source_doc_number_index", "extension_counts",
            "source_extension_counts", "name_order", "size_order", "source_name_order",
            "source_size_order", "word_index", "latin_all_index", "latin_file_index",
            "latin_vocabulary_all")}
        try:
            app.records = records
            app.source_counts = {"Alpha": 1, "Zeta": 1}
            app.source_physical_map = layout["physical_map"]
            app.physical_source_counts = {"OCR": 2}
            app.ocr_public_browser = layout["browsers"]
            app.folder_tree_data = layout["trees"]
            app.folder_browser_data = layout["browsers"]
            app.fulltext_databases = None
            app.build_indexes()
            zeta = app.search("", ["Zeta"], page_size=10)
            self.assertEqual(zeta["total"], 1)
            self.assertEqual(zeta["results"][0]["doc_id"], "OCR:1")
            self.assertEqual(zeta["results"][0]["Source"], "Zeta")
            self.assertEqual(zeta["results"][0]["SourceName"], "Zeta")
            self.assertEqual(zeta["results"][0]["Folder"], ["d1"])
            self.assertEqual(zeta["results"][0]["DisplayPath"], "Zeta/d1/A.txt")
            alpha = app.search("", ["Alpha"], page_size=10)
            self.assertEqual(alpha["total"], 1)
            self.assertEqual(alpha["results"][0]["DisplayPath"], "Alpha/B.txt")
            self.assertEqual(app.search("", None, page_size=10)["total"], 2)
            self.assertEqual(app.search("", ["Zeta"], folders=["d1"], page_size=10)["total"], 1)
            self.assertEqual(app.search("", ["Zeta"], folders=["Alpha"], page_size=10)["total"], 0)
        finally:
            for key, value in saved.items():
                setattr(app, key, value)
            app.build_indexes()
            app.api_response_cache.clear()
            app.search_response_cache.clear()

    def test_startup_loads_one_pinned_release_and_serves_plain_txt(self):
        tree, browser = self.root / "tree.gz", self.root / "browser.gz"
        tree.write_bytes(gzip.compress(b'{"CCRD":[]}')); browser.write_bytes(gzip.compress(b'{"CCRD":{}}'))
        saved = {key: getattr(app, key) for key in ("records", "record_map", "record_map_index", "sources", "source_counts",
                 "source_physical_map", "physical_source_counts", "ocr_public_browser",
                 "folder_tree_data", "folder_browser_data", "active_ocr_release", "fulltext_databases")}
        try:
            with patch.dict("os.environ", {"CCRD_INDEX_BUCKET_DIR": str(self.bucket), "CCRD_OCR_RELEASE_PATH": self.manifest_path,
                                          "CCRD_OCR_RELEASE_SHA256": self.manifest_sha}), \
                    patch.object(app, "DATA_PATH", self.base), patch.object(app, "FOLDER_TREE_PATH", tree), \
                    patch.object(app, "FOLDER_BROWSER_PATH", browser), patch.object(app, "BASE_DIR", self.root):
                app.load_data()
            self.assertEqual(app.source_counts, {"CCRD": 1, "Test": 1})
            self.assertEqual(app.source_physical_map["Test"], "OCR")
            self.assertEqual(app.record_map["CCRD:1"]["display_name"], "旧文件")
            candidate = self.load()
            app.fulltext_databases = app.FulltextDatabases(
                self.root / "empty-db", in_memory=True, extra_paths=[candidate["database"]])
            client = TestClient(app.app)
            response = client.get("/txt/OCR/VoiceOfML/Test/原目录/book.txt")
            self.assertEqual(response.status_code, 200)
            self.assertEqual(response.text, self.content)
            self.assertIn("inline", response.headers["content-disposition"])
            self.assertEqual(client.get("/txt/OCR/missing.txt").status_code, 404)
            self.assertEqual(client.get("/api/preview/OCR:1").json()["text"], self.content)
            self.assertEqual(client.get("/api/preview/OCR:1").json()["text_quality"], "partially-reviewed")
            preview = client.get("/api/preview/OCR:1").json()
            self.assertEqual(preview["reader_url"],
                             "https://voiceofml-search.hf.space/static/reader.html"
                             "?id=1qkrjrq5u0zsy&title=book&ext=pdf-pages")
            self.assertEqual(client.get("/api/preview/OCR:1").json()["path"], "Test/原目录/book.txt")
            self.assertEqual(client.get("/api/preview/OCR:1").json()["source"], "Test")
            self.assertEqual([item["slug"] for item in client.get("/api/sources").json()], ["CCRD", "Test"])
            self.assertEqual(json.loads(app.sidebar_payload_for_path(""))["sources"][-1]["slug"], "Test")
            self.assertEqual(app.zip_entries(["OCR:1"])[0][0], "VoiceOfML/Test/原目录/book.txt")
            public = app.trim_record(app.record_map["OCR:1"])
            self.assertEqual(public["Source"], "Test")
            self.assertEqual(public["SourceName"], "Test")
            self.assertEqual(public["Folder"], ["原目录"])
            self.assertEqual(public["DisplayPath"], "Test/原目录/book.txt")
            self.assertTrue(public["TextUrl"].startswith("/txt/OCR/VoiceOfML/Test/"))
            self.assertEqual(public["TextQuality"], "partially-reviewed")
            folders = client.get("/api/folders/Test").json()
            self.assertEqual(folders[0]["path"], "")
            self.assertEqual(folders[0]["children"][0]["path"], "原目录")
            contents = client.get("/api/folders/Test/contents", params={"path": "原目录"}).json()
            self.assertEqual(contents["files"][0]["doc_id"], "OCR:1")
            search = client.post("/api/search/Test", json={"q": "手机", "page": 1, "page_size": 10,
                                                           "exact": True, "fulltext": True}).json()
            self.assertEqual(search["total"], 1)
            self.assertEqual(search["results"][0]["Source"], "Test")
            self.assertEqual(search["results"][0]["DisplayPath"], "Test/原目录/book.txt")
            self.assertEqual(client.get("/api/random", params={"sources": "Test"}).json()["Source"], "Test")
            self.assertEqual(client.get("/api/extensions", params={"source": "Test"}).json(),
                             [{"name": "txt", "count": 1}])
            self.assertEqual(client.get("/data/sidebar/global.json").json()["sources"][-1]["slug"], "Test")
            sidebar = client.get("/data/sidebar/sources/Test.json").json()
            self.assertEqual(sidebar["source"], "Test")
            self.assertEqual(sidebar["folders"][0]["path"], "原目录")
            initial = client.get("/data/initial/sources/Test.json").json()
            self.assertEqual(initial["total"], 1)
            self.assertEqual(initial["source"], "Test")
            self.assertIn("/data/initial/sources/Test.json", client.get("/data/initial/manifest.json").json()["urls"])
            self.assertIn("/data/sidebar/sources/Test.json", client.get("/data/sidebar/manifest.json").json()["urls"])
        finally:
            if app.fulltext_databases is not saved["fulltext_databases"] and app.fulltext_databases is not None:
                app.fulltext_databases.close()
            for key, value in saved.items(): setattr(app, key, value)
            app.build_indexes()
            app.api_response_cache.clear()
            app.search_response_cache.clear()


if __name__ == "__main__":
    unittest.main()