wehub-resource-sync
342 lines
9.6 KiB
Go
342 lines
9.6 KiB
Go
package server
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/binary"
|
|
"errors"
|
|
"net/http"
|
|
"os"
|
|
"slices"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
fsgguf "github.com/ollama/ollama/fs/gguf"
|
|
"github.com/ollama/ollama/manifest"
|
|
"github.com/ollama/ollama/types/model"
|
|
)
|
|
|
|
func TestModelListCacheHydratesSummary(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
setTestHome(t, t.TempDir())
|
|
createListCacheModel(t, "list-cache", map[string]any{
|
|
"test.context_length": uint32(4096),
|
|
"test.embedding_length": uint32(384),
|
|
}, "{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}")
|
|
|
|
cache := newModelListCache()
|
|
if err := cache.hydrate(context.Background()); err != nil {
|
|
t.Fatalf("hydrate failed: %v", err)
|
|
}
|
|
|
|
summary, ok := cache.Get(model.ParseName("list-cache"))
|
|
if !ok {
|
|
t.Fatal("list summary missing")
|
|
}
|
|
|
|
if summary.Model != "list-cache:latest" || summary.Name != "list-cache:latest" {
|
|
t.Fatalf("summary model/name = %q/%q, want list-cache:latest", summary.Model, summary.Name)
|
|
}
|
|
if summary.Digest == "" {
|
|
t.Fatal("summary digest is empty")
|
|
}
|
|
if summary.Size == 0 {
|
|
t.Fatal("summary size is zero")
|
|
}
|
|
if summary.Details.Family != "test" || summary.Details.Format != "gguf" {
|
|
t.Fatalf("summary details = %+v, want gguf/test", summary.Details)
|
|
}
|
|
if summary.Details.ContextLength != 4096 {
|
|
t.Fatalf("context length = %d, want 4096", summary.Details.ContextLength)
|
|
}
|
|
if summary.Details.EmbeddingLength != 384 {
|
|
t.Fatalf("embedding length = %d, want 384", summary.Details.EmbeddingLength)
|
|
}
|
|
|
|
for _, capability := range []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityInsert} {
|
|
if !slices.Contains(summary.Capabilities, capability) {
|
|
t.Fatalf("capabilities = %v, want %s", summary.Capabilities, capability)
|
|
}
|
|
}
|
|
|
|
listModel := summary.ListModelResponse()
|
|
if !slices.Contains(listModel.Capabilities, model.CapabilityTools) ||
|
|
listModel.Details.ContextLength != 4096 ||
|
|
listModel.Details.EmbeddingLength != 384 {
|
|
t.Fatalf("list response = %+v, want capabilities/context/embedding", listModel)
|
|
}
|
|
}
|
|
|
|
func TestModelListCacheRefreshUpdatesEntry(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
setTestHome(t, t.TempDir())
|
|
createListCacheModel(t, "list-refresh", map[string]any{"test.context_length": uint32(1024)}, "")
|
|
|
|
cache := newModelListCache()
|
|
if err := cache.hydrate(context.Background()); err != nil {
|
|
t.Fatalf("hydrate failed: %v", err)
|
|
}
|
|
|
|
name := model.ParseName("list-refresh")
|
|
first, ok := cache.Get(name)
|
|
if !ok {
|
|
t.Fatal("list summary missing")
|
|
}
|
|
|
|
changeShowCacheManifest(t, "list-refresh")
|
|
if err := cache.RefreshModel(name); err != nil {
|
|
t.Fatalf("refresh failed: %v", err)
|
|
}
|
|
|
|
refreshed, ok := cache.Get(name)
|
|
if !ok {
|
|
t.Fatal("refreshed list summary missing")
|
|
}
|
|
if refreshed.Digest == first.Digest {
|
|
t.Fatalf("digest did not change after refresh: %s", refreshed.Digest)
|
|
}
|
|
if cache.Len() != 1 {
|
|
t.Fatalf("cache entries = %d, want 1", cache.Len())
|
|
}
|
|
}
|
|
|
|
func TestModelListCacheMutationHooks(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
setTestHome(t, t.TempDir())
|
|
|
|
cache := newModelListCache()
|
|
s := Server{modelCaches: &modelCaches{modelList: cache}}
|
|
|
|
_, digest := createBinFile(t, map[string]any{"test.context_length": uint32(2048)}, nil)
|
|
w := createRequest(t, s.CreateHandler, api.CreateRequest{
|
|
Model: "list-hooks",
|
|
Files: map[string]string{"model.gguf": digest},
|
|
Stream: &stream,
|
|
})
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("create model status = %d, want 200: %s", w.Code, w.Body.String())
|
|
}
|
|
|
|
if _, ok := cache.Get(model.ParseName("list-hooks")); !ok {
|
|
t.Fatal("create did not refresh model list cache")
|
|
}
|
|
|
|
w = createRequest(t, s.CopyHandler, api.CopyRequest{
|
|
Source: "list-hooks",
|
|
Destination: "list-hooks-copy",
|
|
})
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("copy model status = %d, want 200: %s", w.Code, w.Body.String())
|
|
}
|
|
if _, ok := cache.Get(model.ParseName("list-hooks-copy")); !ok {
|
|
t.Fatal("copy did not refresh model list cache")
|
|
}
|
|
|
|
w = createRequest(t, s.DeleteHandler, api.DeleteRequest{Model: "list-hooks-copy"})
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("delete model status = %d, want 200: %s", w.Code, w.Body.String())
|
|
}
|
|
if _, ok := cache.Get(model.ParseName("list-hooks-copy")); ok {
|
|
t.Fatal("delete did not remove model list cache entry")
|
|
}
|
|
}
|
|
|
|
func TestModelListCacheSyncsManifestChanges(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
setTestHome(t, t.TempDir())
|
|
createListCacheModel(t, "list-sync-a", map[string]any{"test.context_length": uint32(1024)}, "")
|
|
|
|
cache := newModelListCache()
|
|
cache.Start(context.Background())
|
|
if err := cache.Wait(context.Background()); err != nil {
|
|
t.Fatalf("wait failed: %v", err)
|
|
}
|
|
|
|
createListCacheModel(t, "list-sync-b", map[string]any{"test.context_length": uint32(2048)}, "")
|
|
models, err := cache.List(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("list failed: %v", err)
|
|
}
|
|
|
|
names := make([]string, 0, len(models))
|
|
for _, m := range models {
|
|
names = append(names, m.Name)
|
|
}
|
|
for _, want := range []string{"list-sync-a:latest", "list-sync-b:latest"} {
|
|
if !slices.Contains(names, want) {
|
|
t.Fatalf("names = %v, want %s", names, want)
|
|
}
|
|
}
|
|
|
|
var other Server
|
|
w := createRequest(t, other.DeleteHandler, api.DeleteRequest{Model: "list-sync-a"})
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("delete model status = %d, want 200: %s", w.Code, w.Body.String())
|
|
}
|
|
|
|
models, err = cache.List(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("list after delete failed: %v", err)
|
|
}
|
|
names = names[:0]
|
|
for _, m := range models {
|
|
names = append(names, m.Name)
|
|
}
|
|
if slices.Contains(names, "list-sync-a:latest") || !slices.Contains(names, "list-sync-b:latest") {
|
|
t.Fatalf("names after delete = %v, want only list-sync-b", names)
|
|
}
|
|
}
|
|
|
|
func TestModelListCacheSyncDropsStaleEntryOnRefreshFailure(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
setTestHome(t, t.TempDir())
|
|
createListCacheModel(t, "list-stale", map[string]any{"test.context_length": uint32(1024)}, "")
|
|
|
|
cache := newModelListCache()
|
|
cache.Start(context.Background())
|
|
if err := cache.Wait(context.Background()); err != nil {
|
|
t.Fatalf("wait failed: %v", err)
|
|
}
|
|
|
|
name := model.ParseName("list-stale")
|
|
if _, ok := cache.Get(name); !ok {
|
|
t.Fatal("list summary missing")
|
|
}
|
|
|
|
changeShowCacheManifest(t, "list-stale")
|
|
cache.build = func(model.Name, *manifest.Manifest) (modelListSummary, error) {
|
|
return modelListSummary{}, errors.New("refresh failed")
|
|
}
|
|
|
|
models, err := cache.List(context.Background())
|
|
if err != nil {
|
|
t.Fatalf("list failed: %v", err)
|
|
}
|
|
if len(models) != 0 {
|
|
t.Fatalf("models = %+v, want stale entry removed", models)
|
|
}
|
|
if _, ok := cache.Get(name); ok {
|
|
t.Fatal("stale entry remained in cache after refresh failure")
|
|
}
|
|
}
|
|
|
|
func TestReadModelListGGUFRejectsMalformedMetadata(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
data []byte
|
|
want string
|
|
}{
|
|
{
|
|
name: "oversized key string",
|
|
data: modelListGGUFTestFile(func(b *bytes.Buffer) {
|
|
writeModelListGGUFHeader(t, b, 1)
|
|
writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1)
|
|
}),
|
|
want: "string",
|
|
},
|
|
{
|
|
name: "oversized skipped string",
|
|
data: modelListGGUFTestFile(func(b *bytes.Buffer) {
|
|
writeModelListGGUFHeader(t, b, 1)
|
|
writeModelListGGUFString(t, b, "unused")
|
|
writeModelListGGUFUint32(t, b, modelListGGUFTypeString)
|
|
writeModelListGGUFUint64(t, b, fsgguf.MaxStringLength+1)
|
|
}),
|
|
want: "string",
|
|
},
|
|
{
|
|
name: "oversized skipped array",
|
|
data: modelListGGUFTestFile(func(b *bytes.Buffer) {
|
|
writeModelListGGUFHeader(t, b, 1)
|
|
writeModelListGGUFString(t, b, "unused")
|
|
writeModelListGGUFUint32(t, b, modelListGGUFTypeArray)
|
|
writeModelListGGUFUint32(t, b, modelListGGUFTypeUint8)
|
|
writeModelListGGUFUint64(t, b, fsgguf.MaxArraySize+1)
|
|
}),
|
|
want: "array size",
|
|
},
|
|
}
|
|
|
|
for _, tt := range cases {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
t.Fatalf("readModelListGGUF panicked: %v", r)
|
|
}
|
|
}()
|
|
|
|
path := t.TempDir() + "/model.gguf"
|
|
if err := os.WriteFile(path, tt.data, 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
_, err := readModelListGGUF(path)
|
|
if err == nil {
|
|
t.Fatal("expected error")
|
|
}
|
|
if !strings.Contains(err.Error(), tt.want) {
|
|
t.Fatalf("error = %v, want substring %q", err, tt.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func createListCacheModel(t *testing.T, name string, kv map[string]any, tmpl string) {
|
|
t.Helper()
|
|
_, digest := createBinFile(t, kv, nil)
|
|
|
|
req := api.CreateRequest{
|
|
Model: name,
|
|
Files: map[string]string{"model.gguf": digest},
|
|
Stream: &stream,
|
|
}
|
|
if tmpl != "" {
|
|
req.Template = tmpl
|
|
}
|
|
|
|
var s Server
|
|
w := createRequest(t, s.CreateHandler, req)
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("create model status = %d, want 200: %s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
|
|
func modelListGGUFTestFile(fn func(*bytes.Buffer)) []byte {
|
|
var b bytes.Buffer
|
|
fn(&b)
|
|
return b.Bytes()
|
|
}
|
|
|
|
func writeModelListGGUFHeader(t *testing.T, b *bytes.Buffer, numKV uint64) {
|
|
t.Helper()
|
|
writeModelListGGUFUint32(t, b, modelListGGUFMagicLE)
|
|
writeModelListGGUFUint32(t, b, 3)
|
|
writeModelListGGUFUint64(t, b, 0)
|
|
writeModelListGGUFUint64(t, b, numKV)
|
|
}
|
|
|
|
func writeModelListGGUFString(t *testing.T, b *bytes.Buffer, s string) {
|
|
t.Helper()
|
|
writeModelListGGUFUint64(t, b, uint64(len(s)))
|
|
if _, err := b.WriteString(s); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func writeModelListGGUFUint32(t *testing.T, b *bytes.Buffer, v uint32) {
|
|
t.Helper()
|
|
if err := binary.Write(b, binary.LittleEndian, v); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func writeModelListGGUFUint64(t *testing.T, b *bytes.Buffer, v uint64) {
|
|
t.Helper()
|
|
if err := binary.Write(b, binary.LittleEndian, v); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|