tests/test_entry.py
1.6 KB · 61 lines · python Raw
1 """Tests for KV cache entry data structures."""
2
3 from __future__ import annotations
4
5 import pytest
6
7 from pqc_kv_cache.entry import EncryptedEntry, EntryMetadata, KVCacheEntry
8
9
10 def test_entry_metadata_roundtrip() -> None:
11 meta = EntryMetadata(
12 tenant_id="t1",
13 session_id="s1",
14 layer_idx=3,
15 position=7,
16 token_id=42,
17 kv_role="both",
18 )
19 d = meta.to_dict()
20 assert d["tenant_id"] == "t1"
21 assert d["layer_idx"] == 3
22 assert d["position"] == 7
23 assert d["token_id"] == 42
24 assert d["kv_role"] == "both"
25
26
27 def test_encrypted_entry_roundtrip() -> None:
28 meta = EntryMetadata(
29 tenant_id="t1", session_id="s1", layer_idx=0, position=0
30 )
31 enc = EncryptedEntry(
32 metadata=meta,
33 nonce="aa" * 12,
34 ciphertext="bb" * 64,
35 key_len=32,
36 sequence_number=9,
37 )
38 data = enc.to_dict()
39 restored = EncryptedEntry.from_dict(data)
40 assert restored.metadata == meta
41 assert restored.nonce == enc.nonce
42 assert restored.ciphertext == enc.ciphertext
43 assert restored.key_len == 32
44 assert restored.sequence_number == 9
45
46
47 def test_kvcache_entry_plaintext_size() -> None:
48 meta = EntryMetadata(tenant_id="t", session_id="s", layer_idx=0, position=0)
49 entry = KVCacheEntry(
50 metadata=meta,
51 key_tensor_bytes=b"\x00" * 64,
52 value_tensor_bytes=b"\x00" * 96,
53 )
54 assert entry.plaintext_size() == 160
55
56
57 def test_entry_metadata_frozen() -> None:
58 meta = EntryMetadata(tenant_id="t", session_id="s", layer_idx=0, position=0)
59 with pytest.raises(Exception):
60 meta.tenant_id = "other" # type: ignore[misc]
61