64 lines
1.8 KiB
Go
64 lines
1.8 KiB
Go
package modelstore
|
|
|
|
import (
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
)
|
|
|
|
func TestListInstalledModelsScansModelsDirWithoutManifest(t *testing.T) {
|
|
dir := t.TempDir()
|
|
store := New(dir, 8)
|
|
if err := os.MkdirAll(store.ModelsDir, 0o755); err != nil {
|
|
t.Fatalf("MkdirAll: %v", err)
|
|
}
|
|
content := []byte("abc")
|
|
modelPath := filepath.Join(store.ModelsDir, "face_det_scrfd_500m_640_rk3588.rknn")
|
|
if err := os.WriteFile(modelPath, content, 0o644); err != nil {
|
|
t.Fatalf("WriteFile: %v", err)
|
|
}
|
|
if err := os.WriteFile(filepath.Join(store.ModelsDir, "readme.txt"), []byte("ignore"), 0o644); err != nil {
|
|
t.Fatalf("WriteFile readme: %v", err)
|
|
}
|
|
|
|
items, err := store.ListInstalledModels()
|
|
if err != nil {
|
|
t.Fatalf("ListInstalledModels: %v", err)
|
|
}
|
|
if len(items) != 1 {
|
|
t.Fatalf("unexpected installed models: %#v", items)
|
|
}
|
|
expectedSHA := sha256.Sum256(content)
|
|
if items[0].Name != "face_det_scrfd_500m_640_rk3588" {
|
|
t.Fatalf("unexpected model name: %#v", items[0])
|
|
}
|
|
if items[0].FileName != "face_det_scrfd_500m_640_rk3588.rknn" {
|
|
t.Fatalf("unexpected file name: %#v", items[0])
|
|
}
|
|
if items[0].Sha256 != hex.EncodeToString(expectedSHA[:]) {
|
|
t.Fatalf("unexpected sha256: %#v", items[0])
|
|
}
|
|
}
|
|
|
|
func TestListInstalledModelsScansUploadedFilesDir(t *testing.T) {
|
|
dir := t.TempDir()
|
|
store := New(dir, 8)
|
|
if err := os.MkdirAll(store.FilesDir(), 0o755); err != nil {
|
|
t.Fatalf("MkdirAll: %v", err)
|
|
}
|
|
modelPath := filepath.Join(store.FilesDir(), "face_det_scrfd_500m_640_rk3588__abc.rknn")
|
|
if err := os.WriteFile(modelPath, []byte("abc"), 0o644); err != nil {
|
|
t.Fatalf("WriteFile: %v", err)
|
|
}
|
|
|
|
items, err := store.ListInstalledModels()
|
|
if err != nil {
|
|
t.Fatalf("ListInstalledModels: %v", err)
|
|
}
|
|
if len(items) != 1 || items[0].FileName == "" {
|
|
t.Fatalf("unexpected installed models: %#v", items)
|
|
}
|
|
}
|