svfs/tests/test_huffman.py

"""huffman.py 的往返测试。

这个模块是独立的、只有标准库依赖的算法实现，理论上应当全绿；
它在测试里的作用是：当打包/解包链路出问题时，可以快速排除"压缩算法本身写错了"
这一可能，把问题范围锁到 core.py 的集成层。
"""

import random
import struct
import unittest

import helpers  # noqa: F401  -- 把 src 加入 sys.path

from sqlite_vfs.huffman import (
    FREQ_SIZE,
    MAGIC,
    compress_bytes,
    decompress_bytes,
)


class RoundTripTest(unittest.TestCase):
    def assert_roundtrip(self, data, label=""):
        freq_blob, compressed_blob = compress_bytes(data)

        self.assertEqual(
            FREQ_SIZE, len(freq_blob), f"{label}: 频率表必须是 256 个 uint32"
        )
        self.assertEqual(MAGIC, compressed_blob[:4], f"{label}: 魔数不正确")
        self.assertEqual(
            len(data),
            struct.unpack(">Q", compressed_blob[4:12])[0],
            f"{label}: 头部记录的原长度不正确",
        )
        self.assertEqual(
            data, decompress_bytes(freq_blob, compressed_blob), f"{label}: 还原结果不一致"
        )
        return compressed_blob

    def test_empty(self):
        freq_blob, compressed_blob = compress_bytes(b"")
        self.assertEqual(FREQ_SIZE, len(freq_blob))
        self.assertEqual(b"", decompress_bytes(freq_blob, compressed_blob))

    def test_single_byte(self):
        self.assert_roundtrip(b"\x00", "单字节")
        self.assert_roundtrip(b"Z", "单字节")

    def test_single_symbol_repeated(self):
        self.assert_roundtrip(b"\x7f" * 5000, "单符号重复")

    def test_two_symbols(self):
        self.assert_roundtrip(b"ab" * 3000, "双符号")

    def test_all_256_byte_values(self):
        self.assert_roundtrip(bytes(range(256)), "全字节值")
        self.assert_roundtrip(bytes(range(256)) * 40, "全字节值重复")

    def test_incompressible_random(self):
        data = random.Random(1).randbytes(20000)
        compressed_blob = self.assert_roundtrip(data, "随机数据")
        self.assertGreater(
            len(compressed_blob),
            len(data),
            "随机数据的压缩结果应当大于原始数据（这正是 core 会回退到 none 的原因）",
        )

    def test_repetitive_text_compresses(self):
        data = b"the quick brown fox jumps over the lazy dog. " * 2000
        compressed_blob = self.assert_roundtrip(data, "重复文本")
        # 该文本有 20 个不同字符，理论码长 ~4.5 bit/字符，压缩后约为原始的 56%
        self.assertLess(
            len(compressed_blob),
            int(len(data) * 0.7),
            f"高度重复的文本应当被显著压缩，实际 {len(compressed_blob)}/{len(data)}",
        )

    def test_utf8_chinese(self):
        self.assert_roundtrip("中文内容，测试哈夫曼编码是否无损。\n".encode("utf-8") * 200, "中文")

    def test_binary_with_nulls_and_high_bytes(self):
        data = bytes([0, 255, 0, 128, 64, 0, 7]) * 500
        self.assert_roundtrip(data, "二进制含 0x00/0xFF")

    def test_just_above_core_threshold(self):
        # core.add_file 的压缩阈值是 1024 字节，边界两侧都要能正确往返
        self.assert_roundtrip(b"A" * 1024, "1024 字节")
        self.assert_roundtrip(b"A" * 1025, "1025 字节")


class HeaderValidationTest(unittest.TestCase):
    def setUp(self):
        # 刻意用多符号数据：单符号数据在解压时走快捷路径，不读 payload
        self.freq_blob, self.compressed_blob = compress_bytes(
            b"the quick brown fox. " * 500
        )

    def test_wrong_freq_blob_size_rejected(self):
        with self.assertRaises(ValueError) as ctx:
            decompress_bytes(self.freq_blob[:-4], self.compressed_blob)
        self.assertIn("频率表", str(ctx.exception))

    def test_too_short_blob_rejected(self):
        with self.assertRaises(ValueError) as ctx:
            decompress_bytes(self.freq_blob, b"HUF")
        self.assertIn("太短", str(ctx.exception))

    def test_bad_magic_rejected(self):
        broken = b"XXXX" + self.compressed_blob[4:]
        with self.assertRaises(ValueError) as ctx:
            decompress_bytes(self.freq_blob, broken)
        self.assertIn("魔数", str(ctx.exception))

    def test_truncated_payload_rejected(self):
        truncated = self.compressed_blob[:-50]
        with self.assertRaises(ValueError) as ctx:
            decompress_bytes(self.freq_blob, truncated)
        self.assertIn("提前结束", str(ctx.exception))

    def test_single_symbol_payload_is_not_validated(self):
        """记录既有行为（非缺陷）：单符号数据解压时直接按原长度展开，完全不读 payload，
        因此 payload 被截断也检测不出来。"""
        freq_blob, compressed_blob = compress_bytes(b"A" * 10000)
        self.assertEqual(
            b"A" * 10000, decompress_bytes(freq_blob, compressed_blob[:20])
        )


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