package model
|
|
import (
|
"archive/tar"
|
"bytes"
|
"net/http"
|
"net/http/httptest"
|
"os"
|
"path/filepath"
|
"testing"
|
)
|
|
func TestRegistryReturnsDefaultSenseVoice(t *testing.T) {
|
profile := DefaultModelProfile()
|
if profile.ID != DefaultModelID {
|
t.Fatalf("default model id = %q, want %q", profile.ID, DefaultModelID)
|
}
|
if profile.BackendKind != BackendSenseVoice {
|
t.Fatalf("backend kind = %q, want %q", profile.BackendKind, BackendSenseVoice)
|
}
|
if _, err := GetModelProfile("missing"); err == nil {
|
t.Fatal("expected unknown model id to fail")
|
}
|
}
|
|
func TestRegistryReturnsMoonshineEnglish(t *testing.T) {
|
profile, err := GetModelProfile(MoonshineModelID)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if profile.BackendKind != BackendMoonshine {
|
t.Fatalf("backend kind = %q, want %q", profile.BackendKind, BackendMoonshine)
|
}
|
if profile.InstallDirName != "moonshine-en" {
|
t.Fatalf("install dir = %q, want moonshine-en", profile.InstallDirName)
|
}
|
if len(profile.DownloadURLs) == 0 {
|
t.Fatal("expected moonshine download URL")
|
}
|
}
|
|
func TestRegistryReturnsParakeetEnglish(t *testing.T) {
|
profile, err := GetModelProfile(ParakeetModelID)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if profile.BackendKind != BackendNemoTransducer {
|
t.Fatalf("backend kind = %q, want %q", profile.BackendKind, BackendNemoTransducer)
|
}
|
if profile.InstallDirName != "parakeet-en" {
|
t.Fatalf("install dir = %q, want parakeet-en", profile.InstallDirName)
|
}
|
if profile.Tier != "advanced" {
|
t.Fatalf("tier = %q, want advanced", profile.Tier)
|
}
|
if len(profile.DownloadURLs) == 0 {
|
t.Fatal("expected parakeet download URL")
|
}
|
}
|
|
func TestEnglishLanguageDefaultsToMoonshine(t *testing.T) {
|
profile, err := GetLanguageProfile(EnglishLanguageID)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if profile.DefaultModelID != MoonshineModelID {
|
t.Fatalf("English default model = %q, want %q", profile.DefaultModelID, MoonshineModelID)
|
}
|
if len(profile.UpgradeModelIDs) != 1 || profile.UpgradeModelIDs[0] != ParakeetModelID {
|
t.Fatalf("English upgrade models = %v, want [%s]", profile.UpgradeModelIDs, ParakeetModelID)
|
}
|
}
|
|
func TestValidateSenseVoiceAnyOfAndAllOf(t *testing.T) {
|
dir := t.TempDir()
|
writeTestFile(t, filepath.Join(dir, "tokens.txt"), "tokens")
|
writeTestFile(t, filepath.Join(dir, "model.onnx"), "model")
|
|
result := ValidateModelDir(DefaultModelProfile(), dir)
|
if !result.Valid {
|
t.Fatalf("expected valid model dir, missing=%v problems=%v", result.Missing, result.Problems)
|
}
|
if got := filepath.Base(result.Files["model"]); got != "model.onnx" {
|
t.Fatalf("model file = %q, want model.onnx", got)
|
}
|
}
|
|
func TestValidateSenseVoicePrefersInt8(t *testing.T) {
|
dir := t.TempDir()
|
writeTestFile(t, filepath.Join(dir, "tokens.txt"), "tokens")
|
writeTestFile(t, filepath.Join(dir, "model.onnx"), "model")
|
writeTestFile(t, filepath.Join(dir, "model.int8.onnx"), "int8")
|
|
result := ValidateModelDir(DefaultModelProfile(), dir)
|
if !result.Valid {
|
t.Fatalf("expected valid model dir, missing=%v problems=%v", result.Missing, result.Problems)
|
}
|
if got := filepath.Base(result.Files["model"]); got != "model.int8.onnx" {
|
t.Fatalf("model file = %q, want model.int8.onnx", got)
|
}
|
}
|
|
func TestValidateSenseVoiceRejectsMissingTokens(t *testing.T) {
|
dir := t.TempDir()
|
writeTestFile(t, filepath.Join(dir, "model.int8.onnx"), "model")
|
|
result := ValidateModelDir(DefaultModelProfile(), dir)
|
if result.Valid {
|
t.Fatal("expected missing tokens to fail validation")
|
}
|
}
|
|
func TestValidateSenseVoiceRejectsZeroByteFile(t *testing.T) {
|
dir := t.TempDir()
|
writeTestFile(t, filepath.Join(dir, "tokens.txt"), "tokens")
|
writeTestFile(t, filepath.Join(dir, "model.int8.onnx"), "")
|
|
result := ValidateModelDir(DefaultModelProfile(), dir)
|
if result.Valid {
|
t.Fatal("expected zero-byte model file to fail validation")
|
}
|
}
|
|
func TestValidateMoonshineRequiredFiles(t *testing.T) {
|
profile, err := GetModelProfile(MoonshineModelID)
|
if err != nil {
|
t.Fatal(err)
|
}
|
dir := t.TempDir()
|
createValidMoonshine(t, dir)
|
|
result := ValidateModelDir(profile, dir)
|
if !result.Valid {
|
t.Fatalf("expected valid moonshine dir, missing=%v problems=%v", result.Missing, result.Problems)
|
}
|
for _, role := range []string{"preprocessor", "encoder", "uncached_decoder", "cached_decoder", "tokens"} {
|
if result.Files[role] == "" {
|
t.Fatalf("missing resolved file role %q", role)
|
}
|
}
|
}
|
|
func TestValidateMoonshineRejectsMissingFile(t *testing.T) {
|
profile, err := GetModelProfile(MoonshineModelID)
|
if err != nil {
|
t.Fatal(err)
|
}
|
dir := t.TempDir()
|
createValidMoonshine(t, dir)
|
if err := os.Remove(filepath.Join(dir, "cached_decode.int8.onnx")); err != nil {
|
t.Fatal(err)
|
}
|
|
result := ValidateModelDir(profile, dir)
|
if result.Valid {
|
t.Fatal("expected missing cached decoder to fail validation")
|
}
|
}
|
|
func TestValidateParakeetRequiredFiles(t *testing.T) {
|
profile, err := GetModelProfile(ParakeetModelID)
|
if err != nil {
|
t.Fatal(err)
|
}
|
dir := t.TempDir()
|
createValidParakeet(t, dir)
|
|
result := ValidateModelDir(profile, dir)
|
if !result.Valid {
|
t.Fatalf("expected valid parakeet dir, missing=%v problems=%v", result.Missing, result.Problems)
|
}
|
for _, role := range []string{"encoder", "decoder", "joiner", "tokens"} {
|
if result.Files[role] == "" {
|
t.Fatalf("missing resolved file role %q", role)
|
}
|
}
|
}
|
|
func TestValidateParakeetRejectsMissingFile(t *testing.T) {
|
profile, err := GetModelProfile(ParakeetModelID)
|
if err != nil {
|
t.Fatal(err)
|
}
|
dir := t.TempDir()
|
createValidParakeet(t, dir)
|
if err := os.Remove(filepath.Join(dir, "joiner.int8.onnx")); err != nil {
|
t.Fatal(err)
|
}
|
|
result := ValidateModelDir(profile, dir)
|
if result.Valid {
|
t.Fatal("expected missing joiner to fail validation")
|
}
|
}
|
|
func TestResolveModelUsesLegacyDirWithoutState(t *testing.T) {
|
root := t.TempDir()
|
createValidSenseVoice(t, filepath.Join(root, "sensevoice"))
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if !resolved.IsUsable() {
|
t.Fatalf("expected usable model, status=%s missing=%v problems=%v", resolved.Status, resolved.Missing, resolved.Problems)
|
}
|
if resolved.SourceDirKind != "legacy" {
|
t.Fatalf("source dir kind = %q, want legacy", resolved.SourceDirKind)
|
}
|
}
|
|
func TestResolveModelPrefersNewDirWithoutState(t *testing.T) {
|
root := t.TempDir()
|
createValidSenseVoice(t, filepath.Join(root, "sensevoice"))
|
createValidSenseVoice(t, filepath.Join(root, "sensevoice-zh"))
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if resolved.SourceDirKind != "new" {
|
t.Fatalf("source dir kind = %q, want new", resolved.SourceDirKind)
|
}
|
if got := filepath.Base(resolved.RootDir); got != "sensevoice-zh" {
|
t.Fatalf("root dir = %q, want sensevoice-zh", got)
|
}
|
}
|
|
func TestResolveModelInvalidStateFallsBackToLegacy(t *testing.T) {
|
root := t.TempDir()
|
writeTestFile(t, filepath.Join(root, "state.json"), "{invalid json")
|
createValidSenseVoice(t, filepath.Join(root, "sensevoice"))
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if resolved.Status != ModelStateInvalidFallbackFound {
|
t.Fatalf("status = %s, want %s", resolved.Status, ModelStateInvalidFallbackFound)
|
}
|
if resolved.SourceDirKind != "legacy" {
|
t.Fatalf("source dir kind = %q, want legacy", resolved.SourceDirKind)
|
}
|
}
|
|
func TestResolveModelStatePointsToBrokenDirFallsBackToLegacy(t *testing.T) {
|
root := t.TempDir()
|
broken := filepath.Join(root, "broken")
|
if err := os.MkdirAll(broken, 0755); err != nil {
|
t.Fatal(err)
|
}
|
writeTestFile(t, filepath.Join(broken, "tokens.txt"), "tokens")
|
createValidSenseVoice(t, filepath.Join(root, "sensevoice"))
|
|
state := NewInstallState()
|
state.InstalledModels[DefaultModelID] = InstalledModelState{
|
Path: "broken",
|
SourceDirKind: "new",
|
}
|
if err := SaveInstallStateToRoot(root, state); err != nil {
|
t.Fatal(err)
|
}
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if resolved.Status != ModelStateInvalidFallbackFound {
|
t.Fatalf("status = %s, want %s", resolved.Status, ModelStateInvalidFallbackFound)
|
}
|
if resolved.SourceDirKind != "legacy" {
|
t.Fatalf("source dir kind = %q, want legacy", resolved.SourceDirKind)
|
}
|
}
|
|
func TestResolveModelRejectsStatePathOutsideModelsRoot(t *testing.T) {
|
root := t.TempDir()
|
outside := t.TempDir()
|
createValidSenseVoice(t, filepath.Join(outside, "outside-model"))
|
createValidSenseVoice(t, filepath.Join(root, "sensevoice-zh"))
|
|
state := NewInstallState()
|
state.InstalledModels[DefaultModelID] = InstalledModelState{
|
Path: filepath.Join(outside, "outside-model"),
|
SourceDirKind: "new",
|
}
|
if err := SaveInstallStateToRoot(root, state); err != nil {
|
t.Fatal(err)
|
}
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if resolved.Status != ModelStateInvalidFallbackFound {
|
t.Fatalf("status = %s, want %s", resolved.Status, ModelStateInvalidFallbackFound)
|
}
|
if got := filepath.Base(resolved.RootDir); got != "sensevoice-zh" {
|
t.Fatalf("root dir = %q, want sensevoice-zh", got)
|
}
|
}
|
|
func TestResolveModelRejectsStatePathEscapingModelsRoot(t *testing.T) {
|
root := t.TempDir()
|
createValidSenseVoice(t, filepath.Join(root, "sensevoice"))
|
|
state := NewInstallState()
|
state.InstalledModels[DefaultModelID] = InstalledModelState{
|
Path: "../outside-model",
|
SourceDirKind: "new",
|
}
|
if err := SaveInstallStateToRoot(root, state); err != nil {
|
t.Fatal(err)
|
}
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if resolved.Status != ModelStateInvalidFallbackFound {
|
t.Fatalf("status = %s, want %s", resolved.Status, ModelStateInvalidFallbackFound)
|
}
|
if resolved.SourceDirKind != "legacy" {
|
t.Fatalf("source dir kind = %q, want legacy", resolved.SourceDirKind)
|
}
|
}
|
|
func TestResolveModelIncompleteDir(t *testing.T) {
|
root := t.TempDir()
|
dir := filepath.Join(root, "sensevoice-zh")
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
t.Fatal(err)
|
}
|
writeTestFile(t, filepath.Join(dir, "model.int8.onnx"), "model")
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if resolved.Status != ModelInstalledIncomplete {
|
t.Fatalf("status = %s, want %s", resolved.Status, ModelInstalledIncomplete)
|
}
|
}
|
|
func TestInstallStateRoundTrip(t *testing.T) {
|
root := t.TempDir()
|
if err := UpdateInstalledModelState(root, DefaultModelProfile(), "sensevoice-zh", "new"); err != nil {
|
t.Fatal(err)
|
}
|
|
state, err := LoadInstallStateFromRoot(root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if state.SelectedModelID != DefaultModelID {
|
t.Fatalf("selected model = %q, want %q", state.SelectedModelID, DefaultModelID)
|
}
|
installed := state.InstalledModels[DefaultModelID]
|
if installed.Path != "sensevoice-zh" {
|
t.Fatalf("installed path = %q, want sensevoice-zh", installed.Path)
|
}
|
}
|
|
func TestDownloadFailureDoesNotRemoveLegacyModel(t *testing.T) {
|
root := t.TempDir()
|
legacyDir := filepath.Join(root, "sensevoice")
|
createValidSenseVoice(t, legacyDir)
|
|
err := DownloadProfile(DefaultModelProfile(), []string{"://bad-url"}, root, nil)
|
if err == nil {
|
t.Fatal("expected download failure")
|
}
|
|
result := ValidateModelDir(DefaultModelProfile(), legacyDir)
|
if !result.Valid {
|
t.Fatalf("legacy model was damaged, missing=%v problems=%v", result.Missing, result.Problems)
|
}
|
}
|
|
func TestDownloadProfileRemovesStalePartialDownload(t *testing.T) {
|
root := t.TempDir()
|
stalePath := filepath.Join(root, ".downloads", DefaultModelID, "old-run", "model_package")
|
writeTestFile(t, stalePath, "partial")
|
|
err := DownloadProfile(DefaultModelProfile(), []string{"://bad-url"}, root, nil)
|
if err == nil {
|
t.Fatal("expected download failure")
|
}
|
if _, statErr := os.Stat(filepath.Dir(stalePath)); !os.IsNotExist(statErr) {
|
t.Fatalf("stale download run still exists or stat failed: %v", statErr)
|
}
|
}
|
|
func TestDownloadProfileFallsBackWhenPrimaryArchiveInvalid(t *testing.T) {
|
root := t.TempDir()
|
validArchive := tarArchive(t, map[string]string{
|
"sensevoice/tokens.txt": "tokens",
|
"sensevoice/model.int8.onnx": "fallback-model",
|
})
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
switch r.URL.Path {
|
case "/bad":
|
w.WriteHeader(http.StatusOK)
|
w.Write([]byte("not a tar archive"))
|
case "/fallback":
|
w.WriteHeader(http.StatusOK)
|
w.Write(validArchive)
|
default:
|
http.NotFound(w, r)
|
}
|
}))
|
defer server.Close()
|
|
err := DownloadProfile(DefaultModelProfile(), []string{server.URL + "/bad", server.URL + "/fallback"}, root, nil)
|
if err != nil {
|
t.Fatal(err)
|
}
|
|
resolved, err := ResolveModelInRoot(DefaultModelID, root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if !resolved.IsUsable() {
|
t.Fatalf("expected usable downloaded model, status=%s missing=%v problems=%v", resolved.Status, resolved.Missing, resolved.Problems)
|
}
|
if got := filepath.Base(resolved.RootDir); got != "sensevoice-zh" {
|
t.Fatalf("root dir = %q, want sensevoice-zh", got)
|
}
|
}
|
|
func TestDownloadProfileReplacesExistingNewDir(t *testing.T) {
|
root := t.TempDir()
|
finalDir := filepath.Join(root, "sensevoice-zh")
|
createValidSenseVoice(t, finalDir)
|
writeTestFile(t, filepath.Join(finalDir, "model.int8.onnx"), "old-model")
|
|
validArchive := tarArchive(t, map[string]string{
|
"sensevoice/tokens.txt": "tokens",
|
"sensevoice/model.int8.onnx": "new-model",
|
})
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
w.WriteHeader(http.StatusOK)
|
w.Write(validArchive)
|
}))
|
defer server.Close()
|
|
if err := DownloadProfile(DefaultModelProfile(), []string{server.URL}, root, nil); err != nil {
|
t.Fatal(err)
|
}
|
|
data, err := os.ReadFile(filepath.Join(finalDir, "model.int8.onnx"))
|
if err != nil {
|
t.Fatal(err)
|
}
|
if string(data) != "new-model" {
|
t.Fatalf("model contents = %q, want new-model", string(data))
|
}
|
}
|
|
func TestDownloadProfileBadArchiveDoesNotWriteInstalledState(t *testing.T) {
|
root := t.TempDir()
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
w.WriteHeader(http.StatusOK)
|
w.Write([]byte("not a tar archive"))
|
}))
|
defer server.Close()
|
|
if err := DownloadProfile(DefaultModelProfile(), []string{server.URL}, root, nil); err == nil {
|
t.Fatal("expected bad archive to fail")
|
}
|
|
state, err := LoadInstallStateFromRoot(root)
|
if err != nil {
|
t.Fatal(err)
|
}
|
if _, ok := state.InstalledModels[DefaultModelID]; ok {
|
t.Fatal("bad archive should not write installed model state")
|
}
|
}
|
|
func createValidSenseVoice(t *testing.T, dir string) {
|
t.Helper()
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
t.Fatal(err)
|
}
|
writeTestFile(t, filepath.Join(dir, "tokens.txt"), "tokens")
|
writeTestFile(t, filepath.Join(dir, "model.int8.onnx"), "model")
|
}
|
|
func createValidMoonshine(t *testing.T, dir string) {
|
t.Helper()
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
t.Fatal(err)
|
}
|
writeTestFile(t, filepath.Join(dir, "preprocess.onnx"), "preprocessor")
|
writeTestFile(t, filepath.Join(dir, "encode.int8.onnx"), "encoder")
|
writeTestFile(t, filepath.Join(dir, "uncached_decode.int8.onnx"), "uncached")
|
writeTestFile(t, filepath.Join(dir, "cached_decode.int8.onnx"), "cached")
|
writeTestFile(t, filepath.Join(dir, "tokens.txt"), "tokens")
|
}
|
|
func createValidParakeet(t *testing.T, dir string) {
|
t.Helper()
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
t.Fatal(err)
|
}
|
writeTestFile(t, filepath.Join(dir, "encoder.int8.onnx"), "encoder")
|
writeTestFile(t, filepath.Join(dir, "decoder.int8.onnx"), "decoder")
|
writeTestFile(t, filepath.Join(dir, "joiner.int8.onnx"), "joiner")
|
writeTestFile(t, filepath.Join(dir, "tokens.txt"), "tokens")
|
}
|
|
func tarArchive(t *testing.T, files map[string]string) []byte {
|
t.Helper()
|
var buf bytes.Buffer
|
tw := tar.NewWriter(&buf)
|
for name, content := range files {
|
data := []byte(content)
|
header := &tar.Header{
|
Name: name,
|
Mode: 0644,
|
Size: int64(len(data)),
|
}
|
if err := tw.WriteHeader(header); err != nil {
|
t.Fatal(err)
|
}
|
if _, err := tw.Write(data); err != nil {
|
t.Fatal(err)
|
}
|
}
|
if err := tw.Close(); err != nil {
|
t.Fatal(err)
|
}
|
return buf.Bytes()
|
}
|
|
func writeTestFile(t *testing.T, path, content string) {
|
t.Helper()
|
if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
|
t.Fatal(err)
|
}
|
if err := os.WriteFile(path, []byte(content), 0644); err != nil {
|
t.Fatal(err)
|
}
|
}
|
