@@ -0,0 +1,150 @@
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 TestLmstudioEnricher(t *testing.T) {
15+ t.Parallel()
16+
17+ t.Run("populates context window and name from /api/v1/models", func(t *testing.T) {
18+ t.Parallel()
19+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
20+ require.Equal(t, "/api/v1/models", r.URL.Path)
21+ w.Header().Set("Content-Type", "application/json")
22+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
23+ Data: []lmstudioModelEntry{
24+ {
25+ Key: "qwen2.5-7b-instruct",
26+ DisplayName: "Qwen 2.5 7B Instruct",
27+ MaxContextLength: 32768,
28+ },
29+ {
30+ Key: "llama-3.1-8b",
31+ DisplayName: "Llama 3.1 8B",
32+ MaxContextLength: 131072,
33+ },
34+ },
35+ })
36+ }))
37+ defer srv.Close()
38+
39+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
40+ models := []catwalk.Model{
41+ {ID: "qwen2.5-7b-instruct", Name: "qwen2.5-7b-instruct"},
42+ {ID: "llama-3.1-8b", Name: "llama-3.1-8b"},
43+ {ID: "unknown-model", Name: "unknown-model"},
44+ }
45+
46+ e := &lmstudioEnricher{}
47+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
48+ require.NoError(t, err)
49+ require.Len(t, result, 3)
50+ require.Equal(t, int64(32768), result[0].ContextWindow)
51+ require.Equal(t, "Qwen 2.5 7B Instruct", result[0].Name)
52+ require.Equal(t, int64(131072), result[1].ContextWindow)
53+ require.Equal(t, "Llama 3.1 8B", result[1].Name)
54+ require.Equal(t, int64(0), result[2].ContextWindow)
55+ require.Equal(t, "unknown-model", result[2].Name)
56+ })
57+
58+ t.Run("prefers loaded instance context length over model max", func(t *testing.T) {
59+ t.Parallel()
60+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
61+ w.Header().Set("Content-Type", "application/json")
62+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
63+ Data: []lmstudioModelEntry{
64+ {
65+ Key: "m1",
66+ MaxContextLength: 131072,
67+ LoadedInstances: []lmstudioInstance{
68+ {Config: lmstudioInstanceConfig{ContextLength: 8192}},
69+ },
70+ },
71+ },
72+ })
73+ }))
74+ defer srv.Close()
75+
76+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
77+ models := []catwalk.Model{{ID: "m1", Name: "m1"}}
78+
79+ e := &lmstudioEnricher{}
80+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
81+ require.NoError(t, err)
82+ require.Equal(t, int64(8192), result[0].ContextWindow)
83+ })
84+
85+ t.Run("preserves existing non-zero values", func(t *testing.T) {
86+ t.Parallel()
87+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
88+ w.Header().Set("Content-Type", "application/json")
89+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
90+ Data: []lmstudioModelEntry{
91+ {
92+ Key: "m1",
93+ DisplayName: "Should Not Override",
94+ MaxContextLength: 131072,
95+ },
96+ },
97+ })
98+ }))
99+ defer srv.Close()
100+
101+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
102+ models := []catwalk.Model{
103+ {ID: "m1", Name: "My Custom Name", ContextWindow: 65536},
104+ }
105+
106+ e := &lmstudioEnricher{}
107+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
108+ require.NoError(t, err)
109+ require.Equal(t, int64(65536), result[0].ContextWindow)
110+ require.Equal(t, "My Custom Name", result[0].Name)
111+ })
112+
113+ t.Run("returns models unchanged on HTTP error", func(t *testing.T) {
114+ t.Parallel()
115+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
116+ w.WriteHeader(http.StatusServiceUnavailable)
117+ }))
118+ defer srv.Close()
119+
120+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
121+ models := []catwalk.Model{{ID: "m1"}}
122+
123+ e := &lmstudioEnricher{}
124+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
125+ require.NoError(t, err)
126+ require.Len(t, result, 1)
127+ require.Equal(t, int64(0), result[0].ContextWindow)
128+ })
129+
130+ t.Run("does not override custom name with display name", func(t *testing.T) {
131+ t.Parallel()
132+ srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
133+ w.Header().Set("Content-Type", "application/json")
134+ json.NewEncoder(w).Encode(lmstudioModelsResponse{
135+ Data: []lmstudioModelEntry{
136+ {Key: "m1", DisplayName: "API Name"},
137+ },
138+ })
139+ }))
140+ defer srv.Close()
141+
142+ cfg := Config{ID: "test-lmstudio", BaseURL: srv.URL}
143+ models := []catwalk.Model{{ID: "m1", Name: "User Name"}}
144+
145+ e := &lmstudioEnricher{}
146+ result, err := e.EnrichModels(context.Background(), cfg, &mockResolver{}, models)
147+ require.NoError(t, err)
148+ require.Equal(t, "User Name", result[0].Name)
149+ })
150+}