reference/test_kvcache.py

from __future__ import annotations

import json
from pathlib import Path
import tempfile
import unittest

import kvcache


def metadata() -> dict[str, object]:
    return {
        "model_revision": "tiny-lab-model@0123456789abcdef",
        "tokenizer_sha256": "7f" * 32,
        "rope_theta": 10_000.0,
        "layers": 4,
        "kv_heads": 2,
        "head_dim": 8,
        "dtype": "f16",
        "layout": "layer-major-k-then-v",
        "sequence_start": 0,
        "sequence_length": 37,
    }


class CacheTests(unittest.TestCase):
    def test_round_trip_raw_tensor_bytes(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            payload = bytes(range(255))
            kvcache.write_cache(path, metadata(), payload)
            self.assertEqual((metadata(), payload), kvcache.read_cache(path, metadata()))

    def test_inspection_receipt(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            receipt = kvcache.inspect_cache(path)
            self.assertEqual("KVC1", receipt["format"])
            self.assertEqual(2, receipt["payload_bytes"])
            json.dumps(receipt)

    def test_rejects_wrong_model(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            wrong = metadata()
            wrong["model_revision"] = "other@revision"
            with self.assertRaises(kvcache.CacheCompatibilityError):
                kvcache.read_cache(path, wrong)

    def test_rejects_wrong_position(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            wrong = metadata()
            wrong["sequence_start"] = 10
            with self.assertRaises(kvcache.CacheCompatibilityError):
                kvcache.read_cache(path, wrong)

    def test_rejects_corruption(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            raw = bytearray(path.read_bytes())
            raw[-1] ^= 1
            path.write_bytes(raw)
            with self.assertRaises(kvcache.CacheFormatError):
                kvcache.read_cache(path)

    def test_rejects_truncation(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"kv")
            path.write_bytes(path.read_bytes()[:-1])
            with self.assertRaises(kvcache.CacheFormatError):
                kvcache.read_cache(path)

    def test_rejects_incomplete_identity(self) -> None:
        value = metadata()
        value.pop("tokenizer_sha256")
        with tempfile.TemporaryDirectory() as root:
            with self.assertRaises(kvcache.CacheFormatError):
                kvcache.write_cache(Path(root) / "idea.kvc", value, b"kv")

    def test_atomic_replacement(self) -> None:
        with tempfile.TemporaryDirectory() as root:
            path = Path(root) / "idea.kvc"
            kvcache.write_cache(path, metadata(), b"first")
            kvcache.write_cache(path, metadata(), b"second")
            self.assertEqual(b"second", kvcache.read_cache(path)[1])
            self.assertEqual([], list(Path(root).glob("*.tmp")))


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