mirror of
https://github.com/lukaszraczylo/claude-mnemonic.git
synced 2026-10-12 04:55:18 +00:00
lukaszraczylo
a23b8c4936
Explicitly discard best-effort cleanup errors (os.Remove/Close) to satisfy errcheck; add hermetic table-driven tests for checksum mismatch, version-bump cache invalidation, and retry exhaustion.
236 lines
8.6 KiB
Go
236 lines
8.6 KiB
Go
// Package models provides runtime model download and caching.
|
|
package models
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// sha256Hex returns the hex-encoded SHA-256 of b.
|
|
func sha256Hex(t *testing.T, b []byte) string {
|
|
t.Helper()
|
|
h := sha256.Sum256(b)
|
|
return hex.EncodeToString(h[:])
|
|
}
|
|
|
|
// isolateCacheDir points ModelsDir() at a fresh temp directory by overriding
|
|
// the OS cache-dir resolution env vars (HOME and XDG_CACHE_HOME, which
|
|
// os.UserCacheDir honors on Linux and macOS). It also clears
|
|
// CLAUDE_MNEMONIC_MODEL_DIR so EnsureModel takes the download/cache path
|
|
// rather than the local-dir shortcut. It returns the resulting models dir.
|
|
func isolateCacheDir(t *testing.T) string {
|
|
t.Helper()
|
|
home := t.TempDir()
|
|
t.Setenv("HOME", home)
|
|
t.Setenv("XDG_CACHE_HOME", filepath.Join(home, ".cache"))
|
|
t.Setenv("CLAUDE_MNEMONIC_MODEL_DIR", "")
|
|
|
|
dir := ModelsDir()
|
|
require.True(t, strings.HasPrefix(dir, home),
|
|
"ModelsDir() %q must live under temp HOME %q for hermetic test", dir, home)
|
|
return dir
|
|
}
|
|
|
|
// fileServer returns an httptest server that serves body for any request and
|
|
// records how many requests it received.
|
|
func fileServer(t *testing.T, body []byte) (*httptest.Server, *int32) {
|
|
t.Helper()
|
|
var hits int32
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
atomic.AddInt32(&hits, 1)
|
|
_, _ = w.Write(body)
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
return srv, &hits
|
|
}
|
|
|
|
// TestEnsureModel_Download_CachesValidModel verifies the happy path: a model is
|
|
// downloaded, passes checksum verification, and the version file is written.
|
|
func TestEnsureModel_Download_CachesValidModel(t *testing.T) {
|
|
dir := isolateCacheDir(t)
|
|
body := []byte("model-bytes-happy-path")
|
|
srv, hits := fileServer(t, body)
|
|
|
|
const asset = "test-model.onnx"
|
|
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, body))
|
|
require.NoError(t, err)
|
|
|
|
assert.Equal(t, filepath.Join(dir, asset), path)
|
|
got, err := os.ReadFile(path)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, body, got, "cached file content must match downloaded body")
|
|
|
|
ver, err := os.ReadFile(filepath.Join(dir, ".model_version"))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, ModelVersion, string(ver), "version file must record current ModelVersion")
|
|
|
|
assert.Equal(t, int32(1), atomic.LoadInt32(hits), "exactly one download on first fetch")
|
|
}
|
|
|
|
// TestEnsureModel_ChecksumMismatch_FailsAndRemoves verifies that a downloaded
|
|
// file whose content does not match the expected SHA-256 is rejected with a
|
|
// verification error and is not left behind in the cache.
|
|
func TestEnsureModel_ChecksumMismatch_FailsAndRemoves(t *testing.T) {
|
|
dir := isolateCacheDir(t)
|
|
body := []byte("the-actual-downloaded-bytes")
|
|
srv, _ := fileServer(t, body)
|
|
|
|
const asset = "mismatch-model.onnx"
|
|
wrongSum := sha256Hex(t, []byte("a-different-expected-payload"))
|
|
|
|
path, err := EnsureModel(asset, srv.URL, wrongSum)
|
|
require.Error(t, err)
|
|
assert.Empty(t, path)
|
|
assert.Contains(t, err.Error(), "checksum verification")
|
|
|
|
_, statErr := os.Stat(filepath.Join(dir, asset))
|
|
assert.True(t, os.IsNotExist(statErr),
|
|
"corrupt download must be removed from cache, got statErr=%v", statErr)
|
|
}
|
|
|
|
// TestEnsureModel_VersionBump_InvalidatesCache verifies that when the on-disk
|
|
// .model_version differs from ModelVersion, the cache directory is cleared and
|
|
// the model is re-downloaded fresh.
|
|
func TestEnsureModel_VersionBump_InvalidatesCache(t *testing.T) {
|
|
dir := isolateCacheDir(t)
|
|
require.NoError(t, os.MkdirAll(dir, 0700))
|
|
|
|
// Seed a stale cache: an old version file, a leftover model file, and an
|
|
// unrelated sidecar file. Only the version-bump branch wipes the whole dir
|
|
// (os.RemoveAll), so the sidecar's removal uniquely proves invalidation ran
|
|
// — a plain checksum-driven re-download would leave the sidecar in place.
|
|
const asset = "versioned-model.onnx"
|
|
staleModel := filepath.Join(dir, asset)
|
|
staleSidecar := filepath.Join(dir, "leftover-sidecar.bin")
|
|
require.NoError(t, os.WriteFile(staleModel, []byte("stale-garbage"), 0600))
|
|
require.NoError(t, os.WriteFile(staleSidecar, []byte("unrelated-leftover"), 0600))
|
|
require.NoError(t, os.WriteFile(filepath.Join(dir, ".model_version"), []byte("OLD-VERSION"), 0600))
|
|
|
|
body := []byte("freshly-downloaded-model-bytes")
|
|
srv, hits := fileServer(t, body)
|
|
|
|
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, body))
|
|
require.NoError(t, err)
|
|
|
|
got, err := os.ReadFile(path)
|
|
require.NoError(t, err)
|
|
assert.Equal(t, body, got, "stale cached file must be replaced by fresh download")
|
|
|
|
ver, err := os.ReadFile(filepath.Join(dir, ".model_version"))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, ModelVersion, string(ver), "version file must be rewritten to current ModelVersion")
|
|
|
|
assert.Equal(t, int32(1), atomic.LoadInt32(hits),
|
|
"version bump must force exactly one re-download, not serve the stale file")
|
|
|
|
_, sidecarErr := os.Stat(staleSidecar)
|
|
assert.True(t, os.IsNotExist(sidecarErr),
|
|
"version bump must wipe the whole cache dir, removing unrelated sidecar files")
|
|
}
|
|
|
|
// TestEnsureModel_SameVersionValidCache_NoDownload verifies that a matching
|
|
// version file plus a checksum-valid cached file short-circuits the download.
|
|
func TestEnsureModel_SameVersionValidCache_NoDownload(t *testing.T) {
|
|
dir := isolateCacheDir(t)
|
|
require.NoError(t, os.MkdirAll(dir, 0700))
|
|
|
|
const asset = "cached-model.onnx"
|
|
body := []byte("already-cached-valid-bytes")
|
|
require.NoError(t, os.WriteFile(filepath.Join(dir, asset), body, 0600))
|
|
require.NoError(t, os.WriteFile(filepath.Join(dir, ".model_version"), []byte(ModelVersion), 0600))
|
|
|
|
// Server would record a hit if EnsureModel wrongly re-downloaded.
|
|
srv, hits := fileServer(t, []byte("should-not-be-fetched"))
|
|
|
|
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, body))
|
|
require.NoError(t, err)
|
|
assert.Equal(t, filepath.Join(dir, asset), path)
|
|
assert.Equal(t, int32(0), atomic.LoadInt32(hits), "valid same-version cache must not trigger a download")
|
|
}
|
|
|
|
// TestEnsureModel_RetryExhaustion_ReturnsError verifies that when every download
|
|
// attempt fails (server always returns HTTP 500), EnsureModel exhausts its
|
|
// retries and returns an error reporting the attempt count, leaving no model
|
|
// file behind.
|
|
//
|
|
// Note: downloadWithRetries sleeps with exponential backoff between attempts,
|
|
// so this test takes ~6s by design; the backoff is production behavior and is
|
|
// intentionally not stubbed.
|
|
func TestEnsureModel_RetryExhaustion_ReturnsError(t *testing.T) {
|
|
if testing.Short() {
|
|
t.Skip("skipping retry-exhaustion test (real backoff sleeps) in -short mode")
|
|
}
|
|
|
|
dir := isolateCacheDir(t)
|
|
|
|
var hits int32
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
atomic.AddInt32(&hits, 1)
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
}))
|
|
t.Cleanup(srv.Close)
|
|
|
|
const asset = "unreachable-model.onnx"
|
|
path, err := EnsureModel(asset, srv.URL, sha256Hex(t, []byte("never-served")))
|
|
require.Error(t, err)
|
|
assert.Empty(t, path)
|
|
assert.Contains(t, err.Error(), "attempts", "error must report retry exhaustion")
|
|
|
|
assert.Equal(t, int32(maxRetries), atomic.LoadInt32(&hits),
|
|
"server must be hit exactly maxRetries times")
|
|
|
|
_, statErr := os.Stat(filepath.Join(dir, asset))
|
|
assert.True(t, os.IsNotExist(statErr), "failed download must not leave a model file")
|
|
}
|
|
|
|
// TestEnsureModel_LocalDir verifies the CLAUDE_MNEMONIC_MODEL_DIR shortcut:
|
|
// a present, checksum-valid file is returned directly, and a checksum mismatch
|
|
// or missing file is reported without any network access.
|
|
func TestEnsureModel_LocalDir(t *testing.T) {
|
|
body := []byte("local-dir-model-bytes")
|
|
validSum := sha256Hex(t, body)
|
|
|
|
tests := []struct {
|
|
name string
|
|
writeFile bool
|
|
expectSum string
|
|
wantErr string
|
|
}{
|
|
{name: "valid_local_model", writeFile: true, expectSum: validSum, wantErr: ""},
|
|
{name: "checksum_mismatch", writeFile: true, expectSum: sha256Hex(t, []byte("other")), wantErr: "checksum mismatch"},
|
|
{name: "missing_file", writeFile: false, expectSum: validSum, wantErr: "not found"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
localDir := t.TempDir()
|
|
t.Setenv("CLAUDE_MNEMONIC_MODEL_DIR", localDir)
|
|
|
|
const asset = "local-model.onnx"
|
|
if tt.writeFile {
|
|
require.NoError(t, os.WriteFile(filepath.Join(localDir, asset), body, 0600))
|
|
}
|
|
|
|
path, err := EnsureModel(asset, "http://127.0.0.1:0/should-not-be-used", tt.expectSum)
|
|
if tt.wantErr != "" {
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), tt.wantErr)
|
|
assert.Empty(t, path)
|
|
return
|
|
}
|
|
require.NoError(t, err)
|
|
assert.Equal(t, filepath.Join(localDir, asset), path)
|
|
})
|
|
}
|
|
}
|