| 1 | from concurrent.futures import ThreadPoolExecutor |
| 2 | import os |
| 3 | import mmap |
| 4 | from pathlib import Path |
| 5 | import tempfile |
| 6 | import unittest |
| 7 | |
| 8 | from studio_ai_cache import DiskCache |
| 9 | |
| 10 | |
| 11 | class CacheTests(unittest.TestCase): |
| 12 | def setUp(self): |
| 13 | self.temporary = tempfile.TemporaryDirectory() |
| 14 | self.addCleanup(self.temporary.cleanup) |
| 15 | self.root = Path(self.temporary.name) |
| 16 | self.key = os.urandom(32) |
| 17 | self.cache = DiskCache(self.root / "ssd", self.root / "hdd", 2 * mmap.PAGESIZE, 3 * mmap.PAGESIZE, self.key) |
| 18 | self.addCleanup(self.cache.close) |
| 19 | |
| 20 | def test_hdd_hit_promotes_to_ssd_and_both_limits_hold(self): |
| 21 | for name in ("a.bin", "b.bin", "c.bin"): |
| 22 | self.cache.store(name, name.encode() * 100) |
| 23 | target = bytearray(500) |
| 24 | self.assertEqual(self.cache.load("a.bin", target), "hdd") |
| 25 | self.assertEqual(target, b"a.bin" * 100) |
| 26 | self.assertEqual(self.cache.load("a.bin", target), "ssd") |
| 27 | for directory, limit in (("ssd", 2 * mmap.PAGESIZE), ("hdd", 3 * mmap.PAGESIZE)): |
| 28 | self.assertLessEqual(sum(p.stat().st_size for p in (self.root / directory).glob("*.bin")), limit) |
| 29 | |
| 30 | def test_overflow_evicts_cold_blocks(self): |
| 31 | for name in ("a.bin", "b.bin", "c.bin", "d.bin", "e.bin"): |
| 32 | self.cache.store(name, b"fixture") |
| 33 | self.assertFalse(self.cache.acquire("a.bin", "request")) |
| 34 | self.assertFalse(self.cache.acquire("b.bin", "request")) |
| 35 | self.assertTrue(self.cache.acquire("c.bin", "request")) |
| 36 | self.cache.release("request") |
| 37 | |
| 38 | def test_pinned_request_survives_eviction_and_full_cache_skips_store(self): |
| 39 | self.cache.close() |
| 40 | self.cache = DiskCache(self.root / "ssd", self.root / "hdd", mmap.PAGESIZE, mmap.PAGESIZE, self.key) |
| 41 | self.addCleanup(self.cache.close) |
| 42 | self.cache.store("a.bin", b"first") |
| 43 | self.assertTrue(self.cache.acquire("a.bin", "request")) |
| 44 | self.cache.store("b.bin", b"other") |
| 45 | self.assertFalse(self.cache.acquire("b.bin", "other-request")) |
| 46 | target = bytearray(5) |
| 47 | self.cache.load("a.bin", target) |
| 48 | self.assertEqual(target, b"first") |
| 49 | self.cache.release("request") |
| 50 | self.cache.store("b.bin", b"other") |
| 51 | self.assertFalse(self.cache.acquire("a.bin", "request")) |
| 52 | |
| 53 | def test_restart_recovers_encrypted_blocks_with_restricted_permissions(self): |
| 54 | payload = b"synthetic private prompt fixture" * 10 |
| 55 | self.cache.store("namespace/a.bin", payload) |
| 56 | for path in self.root.rglob("*.bin"): |
| 57 | self.assertNotIn(payload, path.read_bytes()) |
| 58 | self.assertEqual(path.stat().st_mode & 0o777, 0o600) |
| 59 | self.cache.close() |
| 60 | restored = DiskCache(self.root / "ssd", self.root / "hdd", 2 * mmap.PAGESIZE, 3 * mmap.PAGESIZE, self.key) |
| 61 | self.addCleanup(restored.close) |
| 62 | target = bytearray(len(payload)) |
| 63 | self.assertEqual(restored.load("namespace/a.bin", target), "ssd") |
| 64 | self.assertEqual(target, payload) |
| 65 | |
| 66 | def test_corrupt_ssd_falls_back_to_hdd_and_repairs_hot_copy(self): |
| 67 | self.cache.store("a.bin", b"authenticated fixture") |
| 68 | path = self.root / "ssd/a.bin" |
| 69 | contents = bytearray(path.read_bytes()) |
| 70 | contents[20] ^= 1 |
| 71 | path.write_bytes(contents) |
| 72 | target = bytearray(21) |
| 73 | self.assertEqual(self.cache.load("a.bin", target), "hdd") |
| 74 | self.assertEqual(target, b"authenticated fixture") |
| 75 | self.assertEqual(self.cache.load("a.bin", target), "ssd") |
| 76 | |
| 77 | def test_corrupt_both_copies_becomes_cache_miss(self): |
| 78 | self.cache.store("a.bin", b"fixture") |
| 79 | for path in self.root.rglob("*.bin"): |
| 80 | path.write_bytes(b"broken") |
| 81 | with self.assertRaises(FileNotFoundError): |
| 82 | self.cache.load("a.bin", bytearray(7)) |
| 83 | self.assertFalse(self.cache.acquire("a.bin", "request")) |
| 84 | |
| 85 | def test_wrong_key_and_different_namespace_cannot_reuse_state(self): |
| 86 | self.cache.store("old/a.bin", b"fixture") |
| 87 | self.assertFalse(self.cache.acquire("new/a.bin", "request")) |
| 88 | self.cache.close() |
| 89 | other = DiskCache(self.root / "ssd", self.root / "hdd", 2 * mmap.PAGESIZE, 3 * mmap.PAGESIZE, os.urandom(32)) |
| 90 | self.addCleanup(other.close) |
| 91 | with self.assertRaises(FileNotFoundError): |
| 92 | other.load("old/a.bin", bytearray(7)) |
| 93 | |
| 94 | def test_directories_are_exclusive(self): |
| 95 | with self.assertRaises(BlockingIOError): |
| 96 | DiskCache(self.root / "ssd", self.root / "hdd", 2 * mmap.PAGESIZE, 3 * mmap.PAGESIZE, self.key) |
| 97 | |
| 98 | def test_concurrent_loads_keep_request_pins_until_release(self): |
| 99 | self.cache.store("a.bin", b"fixture") |
| 100 | self.assertTrue(self.cache.acquire("a.bin", "request")) |
| 101 | |
| 102 | def exercise(index): |
| 103 | self.cache.store(f"{index}.bin", b"other") |
| 104 | target = bytearray(7) |
| 105 | self.cache.load("a.bin", target) |
| 106 | self.assertEqual(target, b"fixture") |
| 107 | |
| 108 | with ThreadPoolExecutor(max_workers=8) as pool: |
| 109 | list(pool.map(exercise, range(24))) |
| 110 | self.cache.release("request") |
| 111 | self.assertEqual(self.cache.pins, {}) |
| 112 | |
| 113 | |
| 114 | if __name__ == "__main__": |
| 115 | unittest.main() |