solpbc.org
personal memory agent
1# SPDX-License-Identifier: AGPL-3.0-only
2# Copyright (c) 2026 sol pbc
3
4from __future__ import annotations
5
6import hashlib
7from pathlib import Path
8from typing import Any
9
10import httpx
11import pytest
12
13from solstone.think.providers import rerank_install
14
15MODEL_BYTES = b"model fixture bytes"
16TOKENIZER_BYTES = b'{"tokenizer":"fixture"}'
17MODEL_PATH = "onnx/model.onnx"
18TOKENIZER_PATH = "tokenizer.json"
19
20
21def _sha256(data: bytes) -> str:
22 return hashlib.sha256(data).hexdigest()
23
24
25def _fixture_bytes() -> dict[str, bytes]:
26 return {
27 MODEL_PATH: MODEL_BYTES,
28 TOKENIZER_PATH: TOKENIZER_BYTES,
29 }
30
31
32def _fixture_spec() -> rerank_install.RerankModelSpec:
33 data = _fixture_bytes()
34 return rerank_install.RerankModelSpec(
35 repo="Xenova/ms-marco-MiniLM-L-6-v2",
36 revision="a09144355adeed5f58c8ed011d209bf8ee5a1fec",
37 files=(
38 rerank_install.RerankFileSpec(
39 path=MODEL_PATH,
40 sha256=_sha256(data[MODEL_PATH]),
41 size_bytes=len(data[MODEL_PATH]),
42 ),
43 rerank_install.RerankFileSpec(
44 path=TOKENIZER_PATH,
45 sha256=_sha256(data[TOKENIZER_PATH]),
46 size_bytes=len(data[TOKENIZER_PATH]),
47 ),
48 ),
49 )
50
51
52def _fake_download_factory(
53 data: dict[str, bytes],
54 calls: list[str] | None = None,
55):
56 def fake_download(
57 _url: str,
58 dest: Path,
59 file_spec: rerank_install.RerankFileSpec,
60 **_kwargs: Any,
61 ) -> None:
62 if calls is not None:
63 calls.append(file_spec.path)
64 dest.parent.mkdir(parents=True, exist_ok=True)
65 dest.write_bytes(data[file_spec.path])
66
67 return fake_download
68
69
70class _FakeResponse:
71 def __init__(self, payload: bytes, headers: dict[str, str] | None = None) -> None:
72 self._payload = payload
73 self.headers = headers or {}
74
75 def raise_for_status(self) -> None:
76 return
77
78 def iter_bytes(self):
79 midpoint = max(1, len(self._payload) // 2)
80 yield self._payload[:midpoint]
81 yield self._payload[midpoint:]
82
83
84class _FakeStream:
85 def __init__(self, response: _FakeResponse) -> None:
86 self.response = response
87
88 def __enter__(self) -> _FakeResponse:
89 return self.response
90
91 def __exit__(self, *_args: object) -> bool:
92 return False
93
94
95def _patch_stream(
96 monkeypatch: pytest.MonkeyPatch,
97 payload: bytes,
98 headers: dict[str, str] | None = None,
99) -> None:
100 def fake_stream(*_args: Any, **_kwargs: Any) -> _FakeStream:
101 return _FakeStream(_FakeResponse(payload, headers=headers))
102
103 monkeypatch.setattr(httpx, "stream", fake_stream)
104
105
106def _paths(
107 spec: rerank_install.RerankModelSpec, journal_path: Path
108) -> tuple[Path, Path, Path]:
109 model = rerank_install.asset_path(
110 spec.files[0], spec=spec, journal_path=journal_path
111 )
112 tokenizer = rerank_install.asset_path(
113 spec.files[1], spec=spec, journal_path=journal_path
114 )
115 sidecar = rerank_install.sidecar_path(spec=spec, journal_path=journal_path)
116 return model, tokenizer, sidecar
117
118
119def test_fetch_if_missing_writes_files_and_sidecar(tmp_path, monkeypatch) -> None:
120 spec = _fixture_spec()
121 calls: list[str] = []
122 monkeypatch.setattr(
123 rerank_install,
124 "_download_file",
125 _fake_download_factory(_fixture_bytes(), calls),
126 )
127
128 record = rerank_install.install_rerank_model(spec=spec, journal_path=tmp_path)
129
130 model, tokenizer, sidecar = _paths(spec, tmp_path)
131 assert model.read_bytes() == MODEL_BYTES
132 assert tokenizer.read_bytes() == TOKENIZER_BYTES
133 assert sidecar.is_file()
134 assert record == rerank_install.RerankInstallRecord.from_json(
135 sidecar.read_text(encoding="utf-8")
136 )
137 assert calls == [MODEL_PATH, TOKENIZER_PATH]
138
139
140def test_present_valid_install_is_noop(tmp_path, monkeypatch) -> None:
141 spec = _fixture_spec()
142 monkeypatch.setattr(
143 rerank_install,
144 "_download_file",
145 _fake_download_factory(_fixture_bytes()),
146 )
147 rerank_install.install_rerank_model(spec=spec, journal_path=tmp_path)
148 monkeypatch.setattr(
149 rerank_install,
150 "_download_file",
151 lambda *_args, **_kwargs: pytest.fail("download should not start"),
152 )
153
154 record = rerank_install.install_rerank_model(spec=spec, journal_path=tmp_path)
155
156 assert record.files == {
157 file_spec.path: file_spec.sha256 for file_spec in spec.files
158 }
159
160
161def test_check_valid_uses_no_network(tmp_path, monkeypatch) -> None:
162 spec = _fixture_spec()
163 monkeypatch.setattr(
164 rerank_install,
165 "_download_file",
166 _fake_download_factory(_fixture_bytes()),
167 )
168 rerank_install.install_rerank_model(spec=spec, journal_path=tmp_path)
169 monkeypatch.setattr(
170 rerank_install,
171 "_download_file",
172 lambda *_args, **_kwargs: pytest.fail("download should not start"),
173 )
174
175 assert (
176 rerank_install.check_rerank_model(spec=spec, journal_path=tmp_path).revision
177 == spec.revision
178 )
179
180
181def test_check_missing_raises_without_download(tmp_path, monkeypatch) -> None:
182 spec = _fixture_spec()
183 monkeypatch.setattr(
184 rerank_install,
185 "_download_file",
186 lambda *_args, **_kwargs: pytest.fail("download should not start"),
187 )
188
189 with pytest.raises(rerank_install.RerankInstallError) as exc_info:
190 rerank_install.check_rerank_model(spec=spec, journal_path=tmp_path)
191
192 assert exc_info.value.reason_code == "sidecar_missing"
193
194
195def test_force_refetches_when_present(tmp_path, monkeypatch) -> None:
196 spec = _fixture_spec()
197 monkeypatch.setattr(
198 rerank_install,
199 "_download_file",
200 _fake_download_factory(_fixture_bytes()),
201 )
202 rerank_install.install_rerank_model(spec=spec, journal_path=tmp_path)
203 calls: list[str] = []
204 monkeypatch.setattr(
205 rerank_install,
206 "_download_file",
207 _fake_download_factory(_fixture_bytes(), calls),
208 )
209
210 rerank_install.install_rerank_model(force=True, spec=spec, journal_path=tmp_path)
211
212 assert calls == [MODEL_PATH, TOKENIZER_PATH]
213
214
215def test_sha256_mismatch_cleans_partial_install(tmp_path, monkeypatch) -> None:
216 spec = _fixture_spec()
217 bad = bytes([MODEL_BYTES[0] ^ 1]) + MODEL_BYTES[1:]
218 _patch_stream(monkeypatch, bad)
219
220 with pytest.raises(rerank_install.RerankInstallError) as exc_info:
221 rerank_install.install_rerank_model(spec=spec, journal_path=tmp_path)
222
223 model, tokenizer, sidecar = _paths(spec, tmp_path)
224 assert exc_info.value.reason_code == "sha256_mismatch"
225 assert not model.exists()
226 assert not model.with_name(f"{model.name}.tmp").exists()
227 assert not tokenizer.exists()
228 assert not sidecar.exists()
229
230
231def test_size_mismatch_cleans_partial_install(tmp_path, monkeypatch) -> None:
232 spec = _fixture_spec()
233 _patch_stream(monkeypatch, b"x", headers={})
234
235 with pytest.raises(rerank_install.RerankInstallError) as exc_info:
236 rerank_install.install_rerank_model(spec=spec, journal_path=tmp_path)
237
238 model, tokenizer, sidecar = _paths(spec, tmp_path)
239 assert exc_info.value.reason_code == "size_mismatch"
240 assert not model.exists()
241 assert not model.with_name(f"{model.name}.tmp").exists()
242 assert not tokenizer.exists()
243 assert not sidecar.exists()
244
245
246def test_download_verifies_before_rename(tmp_path, monkeypatch) -> None:
247 spec = _fixture_spec()
248 file_spec = spec.files[0]
249 bad = bytes([MODEL_BYTES[0] ^ 1]) + MODEL_BYTES[1:]
250 _patch_stream(monkeypatch, bad)
251 dest = tmp_path / "model.onnx"
252
253 with pytest.raises(rerank_install.RerankInstallError) as exc_info:
254 rerank_install._download_file(
255 "https://example.test/model.onnx", dest, file_spec
256 )
257
258 assert exc_info.value.reason_code == "sha256_mismatch"
259 assert not dest.exists()
260 assert not dest.with_name(f"{dest.name}.tmp").exists()
261
262
263def test_invalid_sidecar_shape_fails_check(tmp_path) -> None:
264 spec = _fixture_spec()
265 sidecar = rerank_install.sidecar_path(spec=spec, journal_path=tmp_path)
266 sidecar.parent.mkdir(parents=True, exist_ok=True)
267 sidecar.write_text("{}\n", encoding="utf-8")
268
269 with pytest.raises(rerank_install.RerankInstallError) as exc_info:
270 rerank_install.check_rerank_model(spec=spec, journal_path=tmp_path)
271
272 assert exc_info.value.reason_code == "sidecar_invalid"