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 TestLlamaCPPWatchOutlivesRequestTimeout(t *testing.T) {
36 requestStarted := make(chan struct{})
37 releaseEvent := make(chan struct{})
38 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
39 w.Header().Set("content-type", "text/event-stream")
40 w.WriteHeader(http.StatusOK)
41 w.(http.Flusher).Flush()
42 close(requestStarted)
43 select {
44 case <-releaseEvent:
45 fmt.Fprint(w, "data: {\"model\":\"test\",\"event\":\"download_progress\",\"data\":{}}\n\n")
46 w.(http.Flusher).Flush()
47 case <-r.Context().Done():
48 }
49 }))
50 defer server.Close()
51
52 client, err := NewLlamaCPPClient(server.URL, "")
53 if err != nil {
54 t.Fatal(err)
55 }
56 client.HTTP.Timeout = 20 * time.Millisecond
57 ctx, cancel := context.WithCancel(context.Background())
58 defer cancel()
59 events := make(chan llamaCPPEvent, 1)
60 watchDone := make(chan error, 1)
61 go func() { watchDone <- client.watch(ctx, events) }()
62 <-requestStarted
63
64 select {
65 case err := <-watchDone:
66 t.Fatalf("watch stopped at ordinary request timeout: %v", err)
67 case <-time.After(3 * client.HTTP.Timeout):
68 }
69 close(releaseEvent)
70 select {
71 case event := <-events:
72 if event.Model != "test" || event.Event != "download_progress" {
73 t.Fatalf("event = %#v", event)
74 }
75 case <-time.After(time.Second):
76 t.Fatal("timed out waiting for delayed SSE event")
77 }
78}
79
80func TestLlamaCPPModelAcceptsTopLevelDownloadProgress(t *testing.T) {
81 var model LlamaCPPModel
82 if err := json.Unmarshal([]byte(`{"id":"test","status":{"value":"downloading"},"progress":{"model.gguf":{"done":512,"total":1024}}}`), &model); err != nil {
83 t.Fatal(err)
84 }
85 progress, ok := downloadProgress(model.Progress)
86 if !ok || progress.Ratio != 0.5 {
87 t.Fatalf("progress = %#v, ok = %v", progress, ok)
88 }
89}
90
91func TestLlamaCPPClientLoadAndDownloadProgress(t *testing.T) {
92 var mu sync.Mutex
93 status := "unloaded"
94 events := make(chan string, 8)
95 watchDone := make(chan struct{}, 2)
96 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
97 if got := r.Header.Get("Authorization"); got != "Bearer secret" {
98 t.Errorf("authorization = %q", got)
99 }
100 switch {
101 case r.URL.Path == "/models/sse":
102 defer func() { watchDone <- struct{}{} }()
103 w.Header().Set("content-type", "text/event-stream")
104 flusher := w.(http.Flusher)
105 for {
106 select {
107 case line := <-events:
108 fmt.Fprintf(w, "data: %s\n\n", line)
109 flusher.Flush()
110 mu.Lock()
111 if strings.Contains(line, `"event":"status_change"`) {
112 status = "loaded"
113 }
114 if strings.Contains(line, `"event":"download_finished"`) {
115 status = "unloaded"
116 }
117 mu.Unlock()
118 case <-r.Context().Done():
119 return
120 }
121 }
122 case r.URL.Path == "/models/load":
123 mu.Lock()
124 status = "loading"
125 mu.Unlock()
126 w.Header().Set("content-type", "application/json")
127 fmt.Fprint(w, `{"success":true}`)
128 events <- `{"model":"test","event":"status_change","data":{"status":"loading","progress":{"stages":["text_model"],"current":"text_model","value":0.5}}}`
129 case r.URL.Path == "/models" && r.Method == http.MethodPost:
130 mu.Lock()
131 status = "downloading"
132 mu.Unlock()
133 fmt.Fprint(w, `{"success":true}`)
134 events <- `{"model":"test","event":"download_progress","data":{"progress":{"model.gguf":{"done":512,"total":1024}}}}`
135 events <- `{"model":"test","event":"download_finished","data":{}}`
136 case r.URL.Path == "/models":
137 mu.Lock()
138 current := status
139 mu.Unlock()
140 modelStatus := map[string]any{"value": current}
141 if current == "downloading" {
142 modelStatus["progress"] = map[string]any{"model.gguf": map[string]any{"done": 512, "total": 1024}}
143 }
144 _ = json.NewEncoder(w).Encode(map[string]any{"data": []any{map[string]any{"id": "test", "status": modelStatus}}})
145 default:
146 http.NotFound(w, r)
147 }
148 }))
149 defer server.Close()
150
151 client, err := NewLlamaCPPClient(server.URL, "secret")
152 if err != nil {
153 t.Fatal(err)
154 }
155 ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
156 defer cancel()
157 var loadProgress []LlamaCPPProgress
158 model, err := client.LoadAndWait(ctx, "test", func(progress LlamaCPPProgress) { loadProgress = append(loadProgress, progress) })
159 if err != nil {
160 t.Fatal(err)
161 }
162 if model.Status.Value != "loaded" {
163 t.Fatalf("status = %q", model.Status.Value)
164 }
165 if len(loadProgress) == 0 {
166 t.Fatal("expected load progress updates")
167 }
168 select {
169 case <-watchDone:
170 case <-ctx.Done():
171 t.Fatal("load progress watcher did not stop")
172 }
173 var download []LlamaCPPProgress
174 if _, err := client.DownloadAndWait(ctx, "test", func(progress LlamaCPPProgress) { download = append(download, progress) }); err != nil {
175 t.Fatal(err)
176 }
177 found := false
178 for _, progress := range download {
179 if progress.HasRatio && progress.Ratio == 0.5 && strings.Contains(progress.Detail, "512 B") {
180 found = true
181 }
182 }
183 if !found {
184 t.Fatalf("download progress = %#v", download)
185 }
186}
187
188func TestLlamaCPPClientRemoveUsesDeleteWithModelQuery(t *testing.T) {
189 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
190 if r.Method != http.MethodDelete || r.URL.Path != "/models" {
191 t.Errorf("request = %s %s", r.Method, r.URL.Path)
192 }
193 if got := r.URL.Query().Get("model"); got != "owner/model:Q4_K_M" {
194 t.Errorf("model = %q", got)
195 }
196 if got := r.Header.Get("Authorization"); got != "Bearer secret" {
197 t.Errorf("authorization = %q", got)
198 }
199 fmt.Fprint(w, `{"success":true}`)
200 }))
201 defer server.Close()
202
203 client, err := NewLlamaCPPClient(server.URL, "secret")
204 if err != nil {
205 t.Fatal(err)
206 }
207 if err := client.Remove(context.Background(), "owner/model:Q4_K_M"); err != nil {
208 t.Fatal(err)
209 }
210}
211
212func TestHuggingFaceDetailsCollectsGGUFQuantizations(t *testing.T) {
213 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
214 if r.URL.Path == "/api/models" {
215 if r.URL.Query().Get("filter") != "gguf" || r.URL.Query().Get("search") != "coder" {
216 t.Errorf("query = %s", r.URL.RawQuery)
217 }
218 fmt.Fprint(w, `[{"id":"owner/model-GGUF","downloads":1200}]`)
219 return
220 }
221 if r.URL.Path == "/api/models/owner/model-GGUF" {
222 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}]}`)
223 return
224 }
225 http.NotFound(w, r)
226 }))
227 defer server.Close()
228
229 client := NewHuggingFaceClient("token")
230 client.BaseURL = server.URL
231 models, err := client.Search(context.Background(), "coder")
232 if err != nil || len(models) != 1 {
233 t.Fatalf("search: models=%#v err=%v", models, err)
234 }
235 details, err := client.Details(context.Background(), models[0].ID)
236 if err != nil {
237 t.Fatal(err)
238 }
239 if details.Gated != "manual" || len(details.Quantizations) != 2 {
240 t.Fatalf("details = %#v", details)
241 }
242 if details.Quantizations[0].Name != "Q4_K_M" || details.Quantizations[0].Size != 5000 {
243 t.Fatalf("first quantization = %#v", details.Quantizations[0])
244 }
245}
246
247func TestSetManagedModelsDoesNotReplaceDiscoveredCatalog(t *testing.T) {
248 SetLiveModels([]Model{{Provider: "openrouter", ID: "vendor/live", Source: "live"}})
249 SetManagedModels([]Model{{Provider: LlamaCPPProviderID, ID: "local", Source: "live"}})
250 t.Cleanup(func() {
251 SetLiveModels(nil)
252 SetManagedModels(nil)
253 })
254 if _, err := FindModel("openrouter", "vendor/live"); err != nil {
255 t.Fatal(err)
256 }
257 if _, err := FindModel(LlamaCPPProviderID, "local"); err != nil {
258 t.Fatal(err)
259 }
260}