1package provider
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "net/http"
8 "net/http/httptest"
9 "strings"
10 "sync"
11 "testing"
12 "time"
13)
14
15func TestNormalizeLlamaCPPURL(t *testing.T) {
16 got, err := NormalizeLlamaCPPURL("http://127.0.0.1:8080/v1/")
17 if err != nil {
18 t.Fatal(err)
19 }
20 if got != "http://127.0.0.1:8080" {
21 t.Fatalf("got %q", got)
22 }
23 if _, err := NormalizeLlamaCPPURL("file:///tmp/router"); err == nil {
24 t.Fatal("expected non-http URL to fail")
25 }
26}
27
28func TestLlamaCPPLoadProgressStages(t *testing.T) {
29 progress, ok := loadProgress(json.RawMessage(`{"progress":{"stages":["text_model","mmproj_model"],"current":"text_model","value":0.5}}`))
30 if !ok || !progress.HasRatio || progress.Ratio != 0.25 || progress.Message != "Loading text model" {
31 t.Fatalf("progress = %#v, ok = %v", progress, ok)
32 }
33}
34
35func TestLlamaCPPClientLoadAndDownloadProgress(t *testing.T) {
36 var mu sync.Mutex
37 status := "unloaded"
38 events := make(chan string, 8)
39 watchDone := make(chan struct{}, 2)
40 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
41 if got := r.Header.Get("Authorization"); got != "Bearer secret" {
42 t.Errorf("authorization = %q", got)
43 }
44 switch {
45 case r.URL.Path == "/models/sse":
46 defer func() { watchDone <- struct{}{} }()
47 w.Header().Set("content-type", "text/event-stream")
48 flusher := w.(http.Flusher)
49 for {
50 select {
51 case line := <-events:
52 fmt.Fprintf(w, "data: %s\n\n", line)
53 flusher.Flush()
54 mu.Lock()
55 if strings.Contains(line, `"event":"status_change"`) {
56 status = "loaded"
57 }
58 if strings.Contains(line, `"event":"download_finished"`) {
59 status = "unloaded"
60 }
61 mu.Unlock()
62 case <-r.Context().Done():
63 return
64 }
65 }
66 case r.URL.Path == "/models/load":
67 mu.Lock()
68 status = "loading"
69 mu.Unlock()
70 w.Header().Set("content-type", "application/json")
71 fmt.Fprint(w, `{"success":true}`)
72 events <- `{"model":"test","event":"status_change","data":{"status":"loading","progress":{"stages":["text_model"],"current":"text_model","value":0.5}}}`
73 case r.URL.Path == "/models" && r.Method == http.MethodPost:
74 mu.Lock()
75 status = "downloading"
76 mu.Unlock()
77 fmt.Fprint(w, `{"success":true}`)
78 events <- `{"model":"test","event":"download_progress","data":{"progress":{"model.gguf":{"done":512,"total":1024}}}}`
79 events <- `{"model":"test","event":"download_finished","data":{}}`
80 case r.URL.Path == "/models":
81 mu.Lock()
82 current := status
83 mu.Unlock()
84 modelStatus := map[string]any{"value": current}
85 if current == "downloading" {
86 modelStatus["progress"] = map[string]any{"model.gguf": map[string]any{"done": 512, "total": 1024}}
87 }
88 _ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]any{"id": "test", "status": modelStatus}}})
89 default:
90 http.NotFound(w, r)
91 }
92 }))
93 defer server.Close()
94
95 client, err := NewLlamaCPPClient(server.URL, "secret")
96 if err != nil {
97 t.Fatal(err)
98 }
99 ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
100 defer cancel()
101 var loadProgress []LlamaCPPProgress
102 model, err := client.LoadAndWait(ctx, "test", func(progress LlamaCPPProgress) { loadProgress = append(loadProgress, progress) })
103 if err != nil {
104 t.Fatal(err)
105 }
106 if model.Status.Value != "loaded" {
107 t.Fatalf("status = %q", model.Status.Value)
108 }
109 if len(loadProgress) == 0 {
110 t.Fatal("expected load progress updates")
111 }
112 select {
113 case <-watchDone:
114 case <-ctx.Done():
115 t.Fatal("load progress watcher did not stop")
116 }
117 var download []LlamaCPPProgress
118 if _, err := client.DownloadAndWait(ctx, "test", func(progress LlamaCPPProgress) { download = append(download, progress) }); err != nil {
119 t.Fatal(err)
120 }
121 found := false
122 for _, progress := range download {
123 if progress.HasRatio && progress.Ratio == 0.5 && strings.Contains(progress.Detail, "512 B") {
124 found = true
125 }
126 }
127 if !found {
128 t.Fatalf("download progress = %#v", download)
129 }
130}
131
132func TestLlamaCPPClientRemoveUsesDeleteWithModelQuery(t *testing.T) {
133 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
134 if r.Method != http.MethodDelete || r.URL.Path != "/models" {
135 t.Errorf("request = %s %s", r.Method, r.URL.Path)
136 }
137 if got := r.URL.Query().Get("model"); got != "owner/model:Q4_K_M" {
138 t.Errorf("model = %q", got)
139 }
140 if got := r.Header.Get("Authorization"); got != "Bearer secret" {
141 t.Errorf("authorization = %q", got)
142 }
143 fmt.Fprint(w, `{"success":true}`)
144 }))
145 defer server.Close()
146
147 client, err := NewLlamaCPPClient(server.URL, "secret")
148 if err != nil {
149 t.Fatal(err)
150 }
151 if err := client.Remove(context.Background(), "owner/model:Q4_K_M"); err != nil {
152 t.Fatal(err)
153 }
154}
155
156func TestHuggingFaceDetailsCollectsGGUFQuantizations(t *testing.T) {
157 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
158 if r.URL.Path == "/api/models" {
159 if r.URL.Query().Get("filter") != "gguf" || r.URL.Query().Get("search") != "coder" {
160 t.Errorf("query = %s", r.URL.RawQuery)
161 }
162 fmt.Fprint(w, `[{"id":"owner/model-GGUF","downloads":1200}]`)
163 return
164 }
165 if r.URL.Path == "/api/models/owner/model-GGUF" {
166 fmt.Fprint(w, `{"id":"owner/model-GGUF","gated":"manual","siblings":[{"rfilename":"model-Q5_K_M.gguf","size":6000},{"rfilename":"model-Q4_K_M-00001-of-00002.gguf","size":2000},{"rfilename":"model-Q4_K_M-00002-of-00002.gguf","size":3000},{"rfilename":"mmproj-F16.gguf","size":1000}]}`)
167 return
168 }
169 http.NotFound(w, r)
170 }))
171 defer server.Close()
172
173 client := NewHuggingFaceClient("token")
174 client.BaseURL = server.URL
175 models, err := client.Search(context.Background(), "coder")
176 if err != nil || len(models) != 1 {
177 t.Fatalf("search: models=%#v err=%v", models, err)
178 }
179 details, err := client.Details(context.Background(), models[0].ID)
180 if err != nil {
181 t.Fatal(err)
182 }
183 if details.Gated != "manual" || len(details.Quantizations) != 2 {
184 t.Fatalf("details = %#v", details)
185 }
186 if details.Quantizations[0].Name != "Q4_K_M" || details.Quantizations[0].Size != 5000 {
187 t.Fatalf("first quantization = %#v", details.Quantizations[0])
188 }
189}
190
191func TestSetManagedModelsDoesNotReplaceDiscoveredCatalog(t *testing.T) {
192 SetLiveModels([]Model{{Provider: "openrouter", ID: "vendor/live", Source: "live"}})
193 SetManagedModels([]Model{{Provider: LlamaCPPProviderID, ID: "local", Source: "live"}})
194 t.Cleanup(func() {
195 SetLiveModels(nil)
196 SetManagedModels(nil)
197 })
198 if _, err := FindModel("openrouter", "vendor/live"); err != nil {
199 t.Fatal(err)
200 }
201 if _, err := FindModel(LlamaCPPProviderID, "local"); err != nil {
202 t.Fatal(err)
203 }
204}