@@ -0,0 +1,135 @@
1+package discover
2+
3+import (
4+ "context"
5+ "encoding/json"
6+ "net/http"
7+ "net/http/httptest"
8+ "testing"
9+
10+ "charm.land/catwalk/pkg/catwalk"
11+ "github.com/stretchr/testify/require"
12+)
13+
14+func TestLlamacppEnricher(t *testing.T) {
15+ t.Parallel()
16+
17+ t.Run("populates context window from meta.n_ctx", func(t *testing.T) {
18+ t.Parallel()
19+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
20+ require.Equal(t, "/v1/models", r.URL.Path)
21+ w.Header().Set("Content-Type", "application/json")
22+ json.NewEncoder(w).Encode(llamacppModelsResponse{
23+ Data: []llamacppModelEntry{
24+ {
25+ ID: "ggml-org/Qwen2.5-Coder-1.5B-Instruct-Q8_0-GGUF",
26+ Meta: llamacppMeta{NCtx: 8192, NCtxTrain: 32768},
27+ },
28+ {
29+ ID: "ggml-org/Llama-3.2-3B-Instruct-Q8_0-GGUF",
30+ Meta: llamacppMeta{NCtx: 4096, NCtxTrain: 131072},
31+ },
32+ },
33+ })
34+ }))
35+ defer srv.Close()
36+
37+ cfg := Config{ID: "test-llamacpp", BaseURL: srv.URL}
38+ models := []catwalk.Model{
39+ {ID: "ggml-org/Qwen2.5-Coder-1.5B-Instruct-Q8_0-GGUF"},
40+ {ID: "ggml-org/Llama-3.2-3B-Instruct-Q8_0-GGUF"},
41+ {ID: "unknown-model"},
42+ }
43+
44+ e := &llamacppEnricher{}
45+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
46+ require.NoError(t, err)
47+ require.Len(t, result, 3)
48+ require.Equal(t, int64(8192), result[0].ContextWindow)
49+ require.Equal(t, int64(4096), result[1].ContextWindow)
50+ require.Equal(t, int64(0), result[2].ContextWindow)
51+ })
52+
53+ t.Run("falls back to n_ctx_train when n_ctx is zero", func(t *testing.T) {
54+ t.Parallel()
55+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
56+ w.Header().Set("Content-Type", "application/json")
57+ json.NewEncoder(w).Encode(llamacppModelsResponse{
58+ Data: []llamacppModelEntry{
59+ {
60+ ID: "m1",
61+ Meta: llamacppMeta{NCtx: 0, NCtxTrain: 131072},
62+ },
63+ },
64+ })
65+ }))
66+ defer srv.Close()
67+
68+ cfg := Config{ID: "test-llamacpp", BaseURL: srv.URL}
69+ models := []catwalk.Model{{ID: "m1"}}
70+
71+ e := &llamacppEnricher{}
72+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
73+ require.NoError(t, err)
74+ require.Equal(t, int64(131072), result[0].ContextWindow)
75+ })
76+
77+ t.Run("preserves existing non-zero context window", func(t *testing.T) {
78+ t.Parallel()
79+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
80+ w.Header().Set("Content-Type", "application/json")
81+ json.NewEncoder(w).Encode(llamacppModelsResponse{
82+ Data: []llamacppModelEntry{
83+ {
84+ ID: "m1",
85+ Meta: llamacppMeta{NCtx: 8192, NCtxTrain: 32768},
86+ },
87+ },
88+ })
89+ }))
90+ defer srv.Close()
91+
92+ cfg := Config{ID: "test-llamacpp", BaseURL: srv.URL}
93+ models := []catwalk.Model{{ID: "m1", ContextWindow: 65536}}
94+
95+ e := &llamacppEnricher{}
96+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
97+ require.NoError(t, err)
98+ require.Equal(t, int64(65536), result[0].ContextWindow)
99+ })
100+
101+ t.Run("returns models unchanged on HTTP error", func(t *testing.T) {
102+ t.Parallel()
103+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
104+ w.WriteHeader(http.StatusServiceUnavailable)
105+ }))
106+ defer srv.Close()
107+
108+ cfg := Config{ID: "test-llamacpp", BaseURL: srv.URL}
109+ models := []catwalk.Model{{ID: "m1"}}
110+
111+ e := &llamacppEnricher{}
112+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
113+ require.NoError(t, err)
114+ require.Len(t, result, 1)
115+ require.Equal(t, int64(0), result[0].ContextWindow)
116+ })
117+
118+ t.Run("returns models unchanged on invalid JSON", func(t *testing.T) {
119+ t.Parallel()
120+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
121+ w.Header().Set("Content-Type", "application/json")
122+ w.Write([]byte("not json"))
123+ }))
124+ defer srv.Close()
125+
126+ cfg := Config{ID: "test-llamacpp", BaseURL: srv.URL}
127+ models := []catwalk.Model{{ID: "m1"}}
128+
129+ e := &llamacppEnricher{}
130+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
131+ require.NoError(t, err)
132+ require.Len(t, result, 1)
133+ require.Equal(t, int64(0), result[0].ContextWindow)
134+ })
135+}