sglang/test/registered/unit/mem_cache/test_cp_l3_posix.py

"""Unit tests for the CP HiCache L3 POSIX disk backend (cp_l3_posix.CpL3DiskSlab).

Stubs the sglang package chain (cross-module import of cp_l3_disk). Uses a real temp FS (O_DIRECT if the
FS supports it, else buffered fallback — both exercised correctly).
"""

import importlib.util
import os
import sys
import tempfile
import types
import unittest
from pathlib import Path

_REPO_ROOT = Path(__file__).resolve().parents[4]
_MEM = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache"


def _load_file(module_name, path):
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


def _load_posix():
    for name, pkg_path in (
        ("sglang", _REPO_ROOT / "python" / "sglang"),
        ("sglang.srt", _REPO_ROOT / "python" / "sglang" / "srt"),
        ("sglang.srt.mem_cache", _MEM),
    ):
        if name not in sys.modules:
            mod = types.ModuleType(name)
            mod.__path__ = [str(pkg_path)]
            sys.modules[name] = mod
    _load_file("sglang.srt.mem_cache.cp_l3_disk", _MEM / "cp_l3_disk.py")
    return _load_file("sglang.srt.mem_cache.cp_l3_posix", _MEM / "cp_l3_posix.py")


px = _load_posix()
PAGE = 8192          # small page for fast tests
NSLOTS = 8


def _h(seed):
    return bytes((seed + i) & 0xFF for i in range(32))


class TestDiskSlab(unittest.TestCase):
    def setUp(self):
        self._td = tempfile.TemporaryDirectory()
        self.slab = px.CpL3DiskSlab(
            os.path.join(self._td.name, "target_kv.slab"), "target_kv", PAGE, NSLOTS
        ).open()

    def tearDown(self):
        self.slab.close()
        self._td.cleanup()

    def test_slot_bytes_aligned(self):
        self.assertEqual(self.slab.slot_bytes % 4096, 0)
        self.assertEqual(self.slab.capacity_bytes, NSLOTS * self.slab.slot_bytes)

    def test_write_read_roundtrip(self):
        payload = bytes((i * 7) & 0xFF for i in range(PAGE))
        h = _h(5)
        crc = self.slab.write_slot(3, h, 78, payload)
        got = self.slab.read_slot(3, expect_hash=h)
        self.assertEqual(got, payload)
        # without expect_hash also intact
        self.assertEqual(self.slab.read_slot(3), payload)

    def test_unwritten_slot_is_miss_not_crash(self):
        self.assertIsNone(self.slab.read_slot(7))  # fresh fallocate = zeros -> bad magic -> None

    def test_wrong_hash_is_miss(self):
        payload = b"\x11" * PAGE
        self.slab.write_slot(2, _h(1), 78, payload)
        self.assertIsNone(self.slab.read_slot(2, expect_hash=_h(99)))

    def test_corruption_detected(self):
        payload = b"\x22" * PAGE
        self.slab.write_slot(1, _h(2), 78, payload)
        # corrupt a payload byte on disk directly
        with open(self.slab.path, "r+b") as f:
            f.seek(1 * self.slab.slot_bytes + 64 + 100)  # into payload
            f.write(b"\xFF")
        self.assertIsNone(self.slab.read_slot(1, expect_hash=_h(2)))  # CRC mismatch -> miss

    def test_read_into_zero_copy(self):
        payload = bytes((i * 3) & 0xFF for i in range(PAGE))
        h = _h(8)
        self.slab.write_slot(4, h, 78, payload)
        buf = self.slab.new_aligned_buffer()
        ok = self.slab.read_into(4, buf, expect_hash=h)
        self.assertTrue(ok)
        self.assertEqual(bytes(buf[64:64 + PAGE]), payload)
        buf.close()

    def test_oob_slot_fails_loud(self):
        with self.assertRaises(ValueError):
            self.slab.write_slot(NSLOTS, _h(0), 1, b"\x00" * PAGE)

    def test_payload_size_mismatch_fails(self):
        with self.assertRaises(ValueError):
            self.slab.write_slot(0, _h(0), 1, b"\x00" * (PAGE - 1))

    def test_fdatasync(self):
        self.slab.write_slot(0, _h(0), 1, b"\x33" * PAGE)
        self.slab.fdatasync()  # must not raise

    def test_read_header_for_cold_rebuild(self):
        # header-only read (cold-rebuild scan): a written slot returns its parsed header; an unwritten slot
        # returns None (fresh fallocate = zeros = bad magic). Cheap: reads one aligned block, not the slot.
        payload = bytes((i * 5) & 0xFF for i in range(PAGE))
        self.slab.write_slot(2, _h(13), 78, payload)
        hdr = self.slab.read_header(2)
        self.assertIsNotNone(hdr)
        self.assertEqual(hdr.payload_kind, "target_kv")
        self.assertEqual(hdr.content_hash, _h(13))  # _h returns 32 raw bytes here
        self.assertEqual(hdr.n_layers, 78)
        self.assertEqual(hdr.page_bytes, PAGE)
        self.assertIsNone(self.slab.read_header(5))  # unwritten -> None


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