@@ -0,0 +1,226 @@ |
| | 1 | +"""Unit tests for :class:`dlm_sway.serve.cache.BackendCache`. |
| | 2 | + |
| | 3 | +The cache is the daemon's heart: an LRU of warm differential backends |
| | 4 | +keyed by the identity tuple over ``ModelSpec``. These tests exercise: |
| | 5 | + |
| | 6 | +* hit / miss / LRU eviction order |
| | 7 | +* concurrent ``get_or_load`` for the same key only loads once |
| | 8 | +* ``close()`` runs on eviction (and tolerates close failures) |
| | 9 | +* ``cache_key_for`` is stable across spec field permutations that |
| | 10 | + don't change identity |
| | 11 | +""" |
| | 12 | + |
| | 13 | +from __future__ import annotations |
| | 14 | + |
| | 15 | +import threading |
| | 16 | +import time |
| | 17 | +from pathlib import Path |
| | 18 | +from typing import Any |
| | 19 | + |
| | 20 | +import pytest |
| | 21 | + |
| | 22 | +from dlm_sway.core.model import ModelSpec |
| | 23 | +from dlm_sway.serve.cache import BackendCache, CachedBackend, cache_key_for |
| | 24 | + |
| | 25 | + |
| | 26 | +class _StubBackend: |
| | 27 | + """Minimal stand-in implementing the loader contract. |
| | 28 | + |
| | 29 | + Tracks how many times :meth:`close` was called so eviction tests |
| | 30 | + can assert the call happened. Doesn't need to satisfy the full |
| | 31 | + ``DifferentialBackend`` Protocol — the cache only ever calls |
| | 32 | + ``.close()`` on it, and tests look up ``.tag`` to identify entries. |
| | 33 | + """ |
| | 34 | + |
| | 35 | + def __init__(self, tag: str) -> None: |
| | 36 | + self.tag = tag |
| | 37 | + self.close_count = 0 |
| | 38 | + |
| | 39 | + def close(self) -> None: |
| | 40 | + self.close_count += 1 |
| | 41 | + |
| | 42 | + |
| | 43 | +def _spec(base: str, *, adapter: Path | None = None) -> ModelSpec: |
| | 44 | + return ModelSpec(base=base, kind="dummy", adapter=adapter) |
| | 45 | + |
| | 46 | + |
| | 47 | +def _seed(cache: BackendCache, spec: ModelSpec, tag: str) -> _StubBackend: |
| | 48 | + """Helper: insert a stub entry under ``spec``'s identity key.""" |
| | 49 | + backend = _StubBackend(tag) |
| | 50 | + key = cache_key_for(spec) |
| | 51 | + entry = CachedBackend( |
| | 52 | + key=key, |
| | 53 | + backend=backend, # type: ignore[arg-type] |
| | 54 | + model_spec=spec, |
| | 55 | + load_seconds=0.0, |
| | 56 | + ) |
| | 57 | + # Use the cache's internal lock + dict directly so we don't trip |
| | 58 | + # the build path. Accessing _entries is fine in unit tests; the |
| | 59 | + # production path goes through get_or_load. |
| | 60 | + with cache._lock: # noqa: SLF001 |
| | 61 | + cache._entries[key] = entry # noqa: SLF001 |
| | 62 | + return backend |
| | 63 | + |
| | 64 | + |
| | 65 | +class TestCacheKey: |
| | 66 | + def test_key_is_stable_across_equivalent_specs(self) -> None: |
| | 67 | + a = _spec("modelA") |
| | 68 | + b = _spec("modelA") |
| | 69 | + assert cache_key_for(a) == cache_key_for(b) |
| | 70 | + |
| | 71 | + def test_key_differs_on_base(self) -> None: |
| | 72 | + assert cache_key_for(_spec("modelA")) != cache_key_for(_spec("modelB")) |
| | 73 | + |
| | 74 | + def test_key_differs_on_adapter(self, tmp_path: Path) -> None: |
| | 75 | + adapter = tmp_path / "ad" |
| | 76 | + adapter.mkdir() |
| | 77 | + with_adapter = _spec("modelA", adapter=adapter) |
| | 78 | + without_adapter = _spec("modelA") |
| | 79 | + assert cache_key_for(with_adapter) != cache_key_for(without_adapter) |
| | 80 | + |
| | 81 | + def test_key_ignores_trust_remote_code(self) -> None: |
| | 82 | + """Two specs differing only in non-identity fields hash equal.""" |
| | 83 | + plain = ModelSpec(base="modelA", kind="dummy") |
| | 84 | + trust = ModelSpec(base="modelA", kind="dummy", trust_remote_code=True) |
| | 85 | + assert cache_key_for(plain) == cache_key_for(trust) |
| | 86 | + |
| | 87 | + |
| | 88 | +class TestCacheLRU: |
| | 89 | + def test_max_size_validation(self) -> None: |
| | 90 | + with pytest.raises(ValueError, match="max_size must be >= 1"): |
| | 91 | + BackendCache(max_size=0) |
| | 92 | + |
| | 93 | + def test_hit_promotes_to_mru(self) -> None: |
| | 94 | + cache = BackendCache(max_size=2) |
| | 95 | + spec_a = _spec("A") |
| | 96 | + spec_b = _spec("B") |
| | 97 | + backend_a = _seed(cache, spec_a, "A") |
| | 98 | + _seed(cache, spec_b, "B") |
| | 99 | + |
| | 100 | + # Touch A so it becomes MRU. get_or_load goes through the |
| | 101 | + # hit path and moves the entry to the end. |
| | 102 | + result = cache.get_or_load(spec_a) |
| | 103 | + assert result.backend is backend_a |
| | 104 | + |
| | 105 | + keys = cache.loaded_keys() |
| | 106 | + assert keys[-1] == cache_key_for(spec_a) |
| | 107 | + assert keys[0] == cache_key_for(spec_b) |
| | 108 | + |
| | 109 | + def test_eviction_closes_lru_backend(self) -> None: |
| | 110 | + """Insert 2 with cap=2, then load a 3rd via get_or_load and |
| | 111 | + confirm the LRU's close() fires.""" |
| | 112 | + cache = BackendCache(max_size=2) |
| | 113 | + spec_a = _spec("A") |
| | 114 | + spec_b = _spec("B") |
| | 115 | + backend_a = _seed(cache, spec_a, "A") |
| | 116 | + _seed(cache, spec_b, "B") |
| | 117 | + |
| | 118 | + # Stub the loader so we don't need a real backend build. |
| | 119 | + third_backend = _StubBackend("C") |
| | 120 | + |
| | 121 | + def _fake_build(spec: ModelSpec, *, adapter_path: Path | None) -> Any: |
| | 122 | + del spec, adapter_path |
| | 123 | + return third_backend |
| | 124 | + |
| | 125 | + import dlm_sway.serve.cache as cache_mod |
| | 126 | + |
| | 127 | + original = cache_mod._build_entry # noqa: SLF001 |
| | 128 | + |
| | 129 | + def _fake_build_entry(spec: ModelSpec, *, key: Any, adapter_path: Any) -> CachedBackend: |
| | 130 | + return CachedBackend( |
| | 131 | + key=key, |
| | 132 | + backend=_fake_build(spec, adapter_path=adapter_path), |
| | 133 | + model_spec=spec, |
| | 134 | + load_seconds=0.0, |
| | 135 | + ) |
| | 136 | + |
| | 137 | + cache_mod._build_entry = _fake_build_entry # type: ignore[assignment] |
| | 138 | + try: |
| | 139 | + cache.get_or_load(_spec("C")) |
| | 140 | + finally: |
| | 141 | + cache_mod._build_entry = original # type: ignore[assignment] |
| | 142 | + |
| | 143 | + assert backend_a.close_count == 1, "LRU eviction should call close()" |
| | 144 | + keys = cache.loaded_keys() |
| | 145 | + assert cache_key_for(spec_a) not in keys |
| | 146 | + assert cache_key_for(spec_b) in keys |
| | 147 | + assert cache_key_for(_spec("C")) in keys |
| | 148 | + |
| | 149 | + def test_evict_all_closes_every_backend(self) -> None: |
| | 150 | + cache = BackendCache(max_size=3) |
| | 151 | + backends = [_seed(cache, _spec(f"M{i}"), f"M{i}") for i in range(3)] |
| | 152 | + cache.evict_all() |
| | 153 | + assert cache.loaded_keys() == [] |
| | 154 | + for b in backends: |
| | 155 | + assert b.close_count == 1 |
| | 156 | + |
| | 157 | + def test_close_failure_is_swallowed(self, caplog: pytest.LogCaptureFixture) -> None: |
| | 158 | + """A backend whose close() raises should not crash the daemon.""" |
| | 159 | + cache = BackendCache(max_size=1) |
| | 160 | + spec = _spec("boom") |
| | 161 | + backend = _seed(cache, spec, "boom") |
| | 162 | + |
| | 163 | + def _raising_close() -> None: |
| | 164 | + raise RuntimeError("close failed") |
| | 165 | + |
| | 166 | + backend.close = _raising_close # type: ignore[method-assign] |
| | 167 | + |
| | 168 | + with caplog.at_level("WARNING"): |
| | 169 | + cache.evict_all() |
| | 170 | + |
| | 171 | + # The error was logged but didn't propagate. |
| | 172 | + assert any("close raised" in r.message for r in caplog.records) |
| | 173 | + assert cache.loaded_keys() == [] |
| | 174 | + |
| | 175 | + |
| | 176 | +class TestSingleFlight: |
| | 177 | + def test_concurrent_get_or_load_loads_once(self) -> None: |
| | 178 | + """Two threads asking for the same spec must result in exactly |
| | 179 | + one underlying build, not two.""" |
| | 180 | + cache = BackendCache(max_size=2) |
| | 181 | + |
| | 182 | + build_count = 0 |
| | 183 | + build_lock = threading.Lock() |
| | 184 | + backend = _StubBackend("solo") |
| | 185 | + |
| | 186 | + import dlm_sway.serve.cache as cache_mod |
| | 187 | + |
| | 188 | + original = cache_mod._build_entry # noqa: SLF001 |
| | 189 | + |
| | 190 | + def _slow_build_entry(spec: ModelSpec, *, key: Any, adapter_path: Any) -> CachedBackend: |
| | 191 | + nonlocal build_count |
| | 192 | + with build_lock: |
| | 193 | + build_count += 1 |
| | 194 | + # Sleep to widen the window for both threads to see a miss. |
| | 195 | + time.sleep(0.05) |
| | 196 | + return CachedBackend( |
| | 197 | + key=key, |
| | 198 | + backend=backend, # type: ignore[arg-type] |
| | 199 | + model_spec=spec, |
| | 200 | + load_seconds=0.0, |
| | 201 | + ) |
| | 202 | + |
| | 203 | + cache_mod._build_entry = _slow_build_entry # type: ignore[assignment] |
| | 204 | + |
| | 205 | + spec = _spec("solo") |
| | 206 | + results: list[CachedBackend] = [] |
| | 207 | + errs: list[BaseException] = [] |
| | 208 | + |
| | 209 | + def _worker() -> None: |
| | 210 | + try: |
| | 211 | + results.append(cache.get_or_load(spec)) |
| | 212 | + except BaseException as exc: # noqa: BLE001 |
| | 213 | + errs.append(exc) |
| | 214 | + |
| | 215 | + try: |
| | 216 | + threads = [threading.Thread(target=_worker) for _ in range(4)] |
| | 217 | + for t in threads: |
| | 218 | + t.start() |
| | 219 | + for t in threads: |
| | 220 | + t.join(timeout=5.0) |
| | 221 | + finally: |
| | 222 | + cache_mod._build_entry = original # type: ignore[assignment] |
| | 223 | + |
| | 224 | + assert errs == [] |
| | 225 | + assert build_count == 1, f"single-flight broken: built {build_count} times" |
| | 226 | + assert all(r.backend is backend for r in results) |