svfs/tests/test_publish.py

"""按内容寻址发布：导出对象 + 清单，客户端按清单取回并校验。

这里刻意用**真实的 HTTP 服务器**跑一遍（而不是只读本地文件），
因为这条路线的卖点正是"普通的静态托管 / CDN 就能用"，
顺带也能量出"取一个文件到底要几次请求、多少字节"。
"""

import functools
import hashlib
import http.server
import json
import os
import shutil
import threading
import time
import unittest

import helpers

from sqlite_vfs import publish
from sqlite_vfs.core import SQLiteVFS

import zlib


class CountingHandler(http.server.SimpleHTTPRequestHandler):
    """记录请求次数与字节数的静态文件服务。"""

    requests = None  # 由测试注入的列表

    def log_message(self, *args):  # 安静点
        pass

    def do_GET(self):
        CountingHandler.requests.append(self.path)
        super().do_GET()

    def send_head(self):
        path = self.translate_path(self.path)
        if os.path.isdir(path):
            # 不列目录，避免噪音
            self.send_error(404, "no listing")
            return None
        return super().send_head()


class HttpServerFixture(helpers.VFSTestCase):
    def serve(self, directory):
        CountingHandler.requests = []
        handler = functools.partial(CountingHandler, directory=directory)
        httpd = http.server.ThreadingHTTPServer(("127.0.0.1", 0), handler)

        def serve():
            try:
                httpd.serve_forever()
            except OSError:
                # 关停时 select 作用在已关闭的 socket 上（Windows 上的收尾竞态），
                # 与被测代码无关，静音以免掩盖真实错误
                pass

        thread = threading.Thread(target=serve, daemon=True)
        thread.start()
        self.addCleanup(httpd.shutdown)
        self.addCleanup(httpd.server_close)
        host, port = httpd.server_address[:2]
        return f"http://{host}:{port}"


class ExportTest(helpers.VFSTestCase):
    def _pack(self, tree=None, compress=True, name="p.svfs"):
        src = helpers.make_tree(
            self.path("src"),
            tree
            or {
                "index.html": b"<div id=app></div>",
                "a.txt": b"same-content",
                "dup.txt": b"same-content",
                "sub/b.js": b"console.log('hello');\n" * 40,
                "empty_dir": None,
            },
        )
        db = self.path(name)
        helpers.pack(src, db, compress=compress, name="demo-pack")
        return src, db

    def test_objects_are_deduplicated(self):
        _src, db = self._pack()
        out = self.path("objects")
        result = publish.export_objects(db, out)
        self.assertEqual(4, result["files"], "4 个文件")
        self.assertEqual(
            3, result["objects"], "a.txt 与 dup.txt 内容相同 → 只导出 3 个对象"
        )
        self.assertGreater(result["dedup_saved"], 0)

    def test_object_content_is_the_original_bytes(self):
        """库里内容是压缩的，导出必须是原始字节——消费方不该需要我们的解压代码。"""
        _src, db = self._pack(compress=True)
        out = self.path("objects")
        publish.export_objects(db, out)

        vfs = SQLiteVFS(db, readonly=True)
        try:
            self.assertEqual("zlib", vfs.get_blob_info("sub/b.js")["compress_algo"])
            digest = vfs.get_blob_info("sub/b.js")["hash"]
        finally:
            vfs.close()

        expected = b"console.log('hello');\n" * 40
        blob = os.path.join(out, "blobs", digest[:2], digest)
        with open(blob, "rb") as f:
            self.assertEqual(expected, f.read())
        self.assertEqual(digest, hashlib.sha256(expected).hexdigest())

    def test_object_name_is_its_hash(self):
        _src, db = self._pack()
        out = self.path("objects")
        publish.export_objects(db, out)
        blob_root = os.path.join(out, "blobs")
        found = 0
        for root, _dirs, files in os.walk(blob_root):
            for name in files:
                with open(os.path.join(root, name), "rb") as f:
                    self.assertEqual(
                        name, hashlib.sha256(f.read()).hexdigest(), f"{name} 名字与内容不符"
                    )
                found += 1
        self.assertEqual(3, found)

    def test_savings_are_reported_in_separate_units(self):
        """去重与压缩必须分开报：把"压缩省下的"算进"去重省下的"会误导选型。"""
        _src, db = self._pack()
        out = self.path("objects")
        result = publish.export_objects(db, out)
        manifest = publish.load_manifest(os.path.join(out, publish.MANIFEST_NAME))

        vfs = SQLiteVFS(db, readonly=True)
        try:
            blobs = vfs.list_blobs()
        finally:
            vfs.close()
        unique = sum(b["size"] for b in blobs)
        stored = sum(b["stored_size"] for b in blobs)

        self.assertEqual(unique, result["unique_bytes"])
        self.assertEqual(stored, result["stored_bytes"])
        self.assertEqual(
            unique, result["objects_bytes"], "导出的对象正好是去重后的原始字节"
        )
        self.assertEqual(
            max(0, manifest["stats"]["total_size"] - unique), result["dedup_saved"]
        )
        self.assertEqual(max(0, unique - stored), result["compression_saved"])
        # 两个口径加起来才是"相对原始文件总量省下的全部"
        self.assertLessEqual(
            result["dedup_saved"] + result["compression_saved"],
            manifest["stats"]["total_size"],
        )

    def test_manifest_contents(self):
        _src, db = self._pack()
        out = self.path("objects")
        publish.export_objects(db, out)
        manifest = publish.load_manifest(os.path.join(out, publish.MANIFEST_NAME))

        self.assertEqual(publish.MANIFEST_FORMAT, manifest["format"])
        self.assertEqual(publish.MANIFEST_VERSION, manifest["version"])
        self.assertEqual("demo-pack", manifest["name"])
        self.assertEqual(1, manifest["snapshot"]["id"])
        self.assertEqual(64, len(manifest["snapshot"]["digest"]))

        entries = manifest["entries"]
        self.assertTrue(entries["empty_dir"]["dir"], "空目录也要出现在清单里")
        self.assertEqual("text/javascript", entries["sub/b.js"]["mime"])
        self.assertEqual(len(b"console.log('hello');\n" * 40), entries["sub/b.js"]["size"])
        self.assertEqual(entries["a.txt"]["hash"], entries["dup.txt"]["hash"])
        self.assertEqual(
            {"files": 4, "directories": 2},
            {k: manifest["stats"][k] for k in ("files", "directories")},
        )

    def test_manifest_is_deterministic_for_the_same_snapshot(self):
        """清单里不放"生成时间"，所以同一快照导出两次得到完全一样的清单（可 diff、可签名）。"""
        _src, db = self._pack()
        first = self.path("out1")
        second = self.path("out2")
        publish.export_objects(db, first)
        publish.export_objects(db, second)
        self.assertEqual(
            helpers.read_bytes(os.path.join(first, publish.MANIFEST_NAME)),
            helpers.read_bytes(os.path.join(second, publish.MANIFEST_NAME)),
        )

    def test_publish_is_read_only_on_the_pack(self):
        _src, db = self._pack()
        before = helpers.md5_of(db)
        publish.export_objects(db, self.path("objects"))
        self.assertEqual(before, helpers.md5_of(db))

    def test_refuses_non_empty_output_without_overwrite(self):
        _src, db = self._pack()
        out = self.path("objects")
        publish.export_objects(db, out)
        with self.assertRaises(FileExistsError):
            publish.export_objects(db, out)
        publish.export_objects(db, out, overwrite=True)

    def test_shard_zero_puts_objects_flat(self):
        _src, db = self._pack()
        out = self.path("objects")
        publish.export_objects(db, out, shard=0)
        names = os.listdir(os.path.join(out, "blobs"))
        self.assertEqual(3, len(names))
        self.assertTrue(all(len(n) == 64 for n in names))

    def test_publish_a_specific_snapshot(self):
        src, db = self._pack(tree={"a.txt": b"v1"})
        vfs = SQLiteVFS(db, compress=True)
        try:
            vfs.create_snapshot("v2", copy_entries=True)
            vfs.remove("a.txt")
            with helpers.quiet():
                vfs.add_file(helpers.write_file(self.path("a2.txt"), b"v2"), "a.txt")
            vfs.recount()
        finally:
            vfs.close()

        out = self.path("objects")
        publish.export_objects(db, out, snapshot_id=1)
        manifest = publish.load_manifest(os.path.join(out, publish.MANIFEST_NAME))
        self.assertEqual(1, manifest["snapshot"]["id"])
        digest = manifest["entries"]["a.txt"]["hash"]
        with open(os.path.join(out, "blobs", digest[:2], digest), "rb") as f:
            self.assertEqual(b"v1", f.read())

    def test_manifest_of_a_nonexistent_snapshot_is_rejected(self):
        _src, db = self._pack()
        with self.assertRaises(FileNotFoundError):
            publish.export_objects(db, self.path("objects"), snapshot_id=99)


class FetchLocalTest(helpers.VFSTestCase):
    def _export(self, tree=None):
        src = helpers.make_tree(
            self.path("src"),
            tree
            or {
                "index.html": b"<div id=app></div>",
                "a.txt": b"same-content",
                "dup.txt": b"same-content",
                "sub/b.js": b"console.log(1)",
                "empty_dir": None,
            },
        )
        db = self.path("p.svfs")
        helpers.pack(src, db, compress=True)
        out = self.path("objects")
        publish.export_objects(db, out)
        return src, out

    def _fetch(self, objects, target, cache=None, **kwargs):
        manifest = publish.load_manifest(os.path.join(objects, publish.MANIFEST_NAME))
        fetcher = publish.ObjectFetcher(
            manifest, local_root=objects, cache_dir=cache, **kwargs
        )
        return fetcher, fetcher.fetch(target, paths=kwargs.pop("paths", None))

    def test_roundtrip_restores_the_tree(self):
        src, objects = self._export()
        target = self.path("restored")
        _fetcher, stats = self._fetch(objects, target)

        helpers.assert_tree_equal(self, helpers.read_tree(src), helpers.read_tree(target))
        self.assertEqual(4, stats.files)
        self.assertEqual(2, stats.directories, "empty_dir 与 sub 都要建出来")
        self.assertEqual(3, stats.objects_downloaded, "一个对象只取一次（两个路径共用）")

    def test_mtime_and_mode_are_preserved(self):
        src, objects = self._export()
        target = self.path("restored")
        self._fetch(objects, target)
        for name in ("a.txt", "sub/b.js"):
            with self.subTest(name=name):
                source = os.path.join(src, *name.split("/"))
                restored = os.path.join(target, *name.split("/"))
                self.assertAlmostEqual(
                    os.path.getmtime(source), os.path.getmtime(restored), delta=2
                )

    def test_partial_fetch_by_file_and_by_prefix(self):
        _src, objects = self._export()
        manifest = publish.load_manifest(os.path.join(objects, publish.MANIFEST_NAME))

        fetcher = publish.ObjectFetcher(manifest, local_root=objects)
        self.assertEqual({"sub/b.js"}, fetcher.resolve(["sub/b.js"]))
        self.assertEqual({"sub/b.js"}, fetcher.resolve(["sub"]))
        self.assertEqual({"sub/b.js"}, fetcher.resolve(["sub/"]))
        self.assertEqual(
            {"a.txt", "sub/b.js"},
            fetcher.resolve(["a.txt", "sub"]),
            "没被选中的同级文件不该被带进来",
        )
        self.assertEqual(
            {"index.html", "a.txt", "dup.txt", "sub/b.js"},
            fetcher.resolve(None),
            "不指定 --path 就是全部文件（不含目录）",
        )
        with self.assertRaises(FileNotFoundError):
            fetcher.resolve(["nope.txt"])

    def test_partial_fetch_writes_only_selected(self):
        _src, objects = self._export()
        target = self.path("partial")
        manifest = publish.load_manifest(os.path.join(objects, publish.MANIFEST_NAME))
        stats = publish.ObjectFetcher(manifest, local_root=objects).fetch(
            target, paths=["sub"]
        )
        self.assertEqual(1, stats.files)
        self.assertEqual(
            {"sub/b.js"}, {k for k, v in helpers.read_tree(target).items() if v is not None}
        )

    def test_cache_is_reused_and_verified(self):
        _src, objects = self._export()
        cache = self.path("cache")
        _fetcher, first = self._fetch(objects, self.path("r1"), cache=cache)
        self.assertEqual(3, first.objects_downloaded)
        self.assertEqual(0, first.objects_from_cache)

        _fetcher, second = self._fetch(objects, self.path("r2"), cache=cache)
        self.assertEqual(0, second.objects_downloaded, "第二次应当全部命中缓存")
        self.assertEqual(3, second.objects_from_cache)
        self.assertEqual(3, second.verified)

    def test_corrupted_cache_entry_is_healed(self):
        _src, objects = self._export()
        cache = self.path("cache")
        self._fetch(objects, self.path("r1"), cache=cache)

        # 把缓存里的某个对象改坏
        blob_root = os.path.join(cache, "blobs")
        victim = None
        for root, _dirs, files in os.walk(blob_root):
            for name in files:
                victim = os.path.join(root, name)
        with open(victim, "wb") as f:
            f.write(b"corrupted!!")

        _fetcher, stats = self._fetch(objects, self.path("r2"), cache=cache)
        self.assertEqual(1, stats.objects_downloaded, "坏掉的缓存应当被重新下载")
        self.assertEqual(2, stats.objects_from_cache)

    def test_link_mode_hardlinks_into_the_cache(self):
        if os.name == "nt" and not hasattr(os, "link"):
            self.skipTest("平台不支持硬链接")
        _src, objects = self._export()
        cache = self.path("cache")
        target = self.path("linked")
        manifest = publish.load_manifest(os.path.join(objects, publish.MANIFEST_NAME))
        publish.ObjectFetcher(
            manifest, local_root=objects, cache_dir=cache, link=True
        ).fetch(target)
        linked = os.path.join(target, "a.txt")
        digest = hashlib.sha256(b"same-content").hexdigest()
        cached = os.path.join(cache, "blobs", digest[:2], digest)
        self.assertEqual(
            os.stat(cached).st_ino, os.stat(linked).st_ino, "应当共用同一个 inode"
        )


class FetchOverHttpTest(HttpServerFixture):
    def _export_and_serve(self):
        src = helpers.make_tree(
            self.path("src"),
            {
                "index.html": b"<div id=app></div>",
                "big.bin": b"x" * (256 * 1024),
                "sub/b.js": b"console.log(1)",
            },
        )
        db = self.path("p.svfs")
        helpers.pack(src, db, compress=True)
        objects = self.path("objects")
        publish.export_objects(db, objects)
        return src, objects, self.serve(objects), db

    def test_fetch_from_a_real_http_server(self):
        src, objects, base, db = self._export_and_serve()
        manifest_url = f"{base}/manifest.json"
        target = self.path("restored")

        manifest = publish.load_manifest(manifest_url)
        fetcher = publish.ObjectFetcher(
            manifest, base_url=base, cache_dir=self.path("cache")
        )
        started = time.time()
        stats = fetcher.fetch(target)
        elapsed = time.time() - started

        helpers.assert_tree_equal(self, helpers.read_tree(src), helpers.read_tree(target))
        self.assertEqual(3, stats.files)
        self.assertEqual(3, stats.objects_downloaded)

        requests = CountingHandler.requests
        manifest_bytes = os.path.getsize(os.path.join(objects, publish.MANIFEST_NAME))
        total_bytes = os.path.getsize(db)
        print(
            f"\n    [HTTP 实测] 请求 {len(requests)} 次"
            f"（清单 1 + 对象 {stats.objects_downloaded}）"
            f"，下载 {stats.bytes_downloaded / 1024:.1f} KB"
            f"（清单 {manifest_bytes / 1024:.1f} KB）"
            f"，耗时 {elapsed * 1000:.0f}ms"
        )
        self.assertEqual(
            1 + stats.objects_downloaded,
            len(requests),
            "总请求数应当是 1 次清单 + 每个对象一次",
        )
        self.assertLess(manifest_bytes, total_bytes, "清单必须远小于整包")

    def test_missing_object_is_reported_clearly(self):
        _src, objects, base, _db = self._export_and_serve()
        manifest = publish.load_manifest(f"{base}/manifest.json")
        # 删掉一个对象，模拟"CDN 上少传了一个文件"
        digest = manifest["entries"]["sub/b.js"]["hash"]
        os.remove(os.path.join(objects, "blobs", digest[:2], digest))

        fetcher = publish.ObjectFetcher(manifest, base_url=base)
        with self.assertRaises(Exception):
            fetcher.fetch(self.path("restored"))

    def test_tampered_object_over_http_is_rejected(self):
        _src, objects, base, _db = self._export_and_serve()
        manifest = publish.load_manifest(f"{base}/manifest.json")
        digest = manifest["entries"]["sub/b.js"]["hash"]
        victim = os.path.join(objects, "blobs", digest[:2], digest)
        with open(victim, "wb") as f:
            f.write(b"console.log(2)")  # 同样长度，但内容不同

        fetcher = publish.ObjectFetcher(manifest, base_url=base, verify=True)
        with self.assertRaises(ValueError) as ctx:
            fetcher.fetch(self.path("restored"))
        self.assertIn("校验失败", str(ctx.exception))

        # 关掉校验就会放过（默认必须是开的）
        manifest2 = publish.load_manifest(f"{base}/manifest.json")
        fetcher2 = publish.ObjectFetcher(manifest2, base_url=base, verify=False)
        fetcher2.fetch(self.path("unverified"))
        with open(os.path.join(self.path("unverified"), "sub", "b.js"), "rb") as f:
            self.assertEqual(b"console.log(2)", f.read())

    def test_server_index_json_gets_actionable_error(self):
        # 最常见的误操作：把服务器根地址（索引 JSON）当清单地址填进来，
        # 报错要直接给出可用的插件 id，而不是一个看不懂的 format=None
        index = self.path("index.json")
        with open(index, "w", encoding="utf-8") as f:
            json.dump({"service": "svfs-plugin-server",
                       "plugins": ["com.svfs.demo", "com.example.other"]}, f)
        with self.assertRaises(ValueError) as ctx:
            publish.load_manifest(index)
        message = str(ctx.exception)
        self.assertIn("服务器索引", message)
        self.assertIn("com.svfs.demo", message)

        empty = self.path("empty-index.json")
        with open(empty, "w", encoding="utf-8") as f:
            json.dump({"service": "svfs-plugin-server", "plugins": []}, f)
        with self.assertRaises(ValueError) as ctx:
            publish.load_manifest(empty)
        self.assertIn("还没有插件", str(ctx.exception))

    def test_manifest_from_url_mismatch_is_reported(self):
        _src, objects, base, _db = self._export_and_serve()
        with self.assertRaises(Exception):
            publish.load_manifest(f"{base}/nope.json")

    def test_bad_manifest_format_is_rejected(self):
        bad = helpers.write_file(
            self.path("bad.json"), json.dumps({"format": "something-else"}).encode()
        )
        with self.assertRaises(ValueError) as ctx:
            publish.load_manifest(bad)
        self.assertIn("svfs-objects", str(ctx.exception))

        old = helpers.write_file(
            self.path("old.json"),
            json.dumps(
                {"format": publish.MANIFEST_FORMAT, "version": 999}
            ).encode(),
        )
        with self.assertRaises(ValueError):
            publish.load_manifest(old)


class PublishCliTest(helpers.VFSTestCase):
    PUBLISH = "sqlite_vfs.cli.publish"
    FETCH = "sqlite_vfs.cli.fetch"

    def _pack(self):
        src = helpers.make_tree(
            self.path("src"), {"a.txt": b"aaa", "sub/b.txt": b"bbb"}
        )
        db = self.path("p.svfs")
        helpers.run_cli("sqlite_vfs.cli.packer", src, db, "--compress")
        return src, db

    def _fetch(self, manifest, target, *extra):
        """CLI 默认用的是用户级缓存目录，测试必须换成临时目录才可复现。"""
        return helpers.run_cli(
            self.FETCH,
            manifest,
            target,
            "--cache",
            self.path("cache"),
            *extra,
        )

    def test_publish_and_fetch_via_cli(self):
        src, db = self._pack()
        objects = self.path("objects")
        published = helpers.run_cli(self.PUBLISH, db, objects)
        self.assertEqual(0, published.returncode, helpers.cli_output(published))
        self.assertIn("发布完成", published.stdout)
        self.assertTrue(os.path.exists(os.path.join(objects, "manifest.json")))

        target = self.path("restored")
        fetched = self._fetch(os.path.join(objects, "manifest.json"), target)
        self.assertEqual(0, fetched.returncode, helpers.cli_output(fetched))
        helpers.assert_tree_equal(self, helpers.read_tree(src), helpers.read_tree(target))

    def test_publish_missing_db_exits_nonzero(self):
        result = helpers.run_cli(self.PUBLISH, self.path("nope.svfs"), self.path("o"))
        self.assertNotEqual(0, result.returncode)
        self.assertNotIn("Traceback", helpers.cli_output(result))

    def test_publish_to_nonempty_dir_requires_overwrite(self):
        _src, db = self._pack()
        objects = self.path("objects")
        helpers.run_cli(self.PUBLISH, db, objects)
        again = helpers.run_cli(self.PUBLISH, db, objects)
        self.assertNotEqual(0, again.returncode)
        self.assertIn("错误", again.stdout)
        forced = helpers.run_cli(self.PUBLISH, db, objects, "--overwrite")
        self.assertEqual(0, forced.returncode, helpers.cli_output(forced))

    def test_manifest_only_mode(self):
        _src, db = self._pack()
        objects = self.path("objects")
        result = helpers.run_cli(self.PUBLISH, db, objects, "--no-objects")
        self.assertEqual(0, result.returncode, helpers.cli_output(result))
        self.assertTrue(os.path.exists(os.path.join(objects, "manifest.json")))
        self.assertFalse(os.path.exists(os.path.join(objects, "blobs")))

    def test_fetch_json_output(self):
        _src, db = self._pack()
        objects = self.path("objects")
        helpers.run_cli(self.PUBLISH, db, objects)
        result = self._fetch(
            os.path.join(objects, "manifest.json"), self.path("restored"), "--json"
        )
        self.assertEqual(0, result.returncode, helpers.cli_output(result))
        stats = json.loads(result.stdout)
        self.assertEqual(2, stats["files"])
        self.assertEqual(2, stats["objects_downloaded"])

    def test_fetch_missing_path_exits_nonzero(self):
        _src, db = self._pack()
        objects = self.path("objects")
        helpers.run_cli(self.PUBLISH, db, objects)
        result = self._fetch(
            os.path.join(objects, "manifest.json"),
            self.path("restored"),
            "--path",
            "nope.txt",
        )
        self.assertNotEqual(0, result.returncode)
        self.assertIn("清单里没有", helpers.cli_output(result))


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