1package ollama
2
3import (
4 "context"
5 "testing"
6 "time"
7)
8
9func TestIsRunning(t *testing.T) {
10 if !IsInstalled() {
11 t.Skip("Ollama is not installed, skipping IsRunning test")
12 }
13
14 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
15 defer cancel()
16
17 running := IsRunning(ctx)
18
19 if running {
20 t.Log("✓ Ollama is running")
21 } else {
22 t.Log("✗ Ollama is not running")
23 }
24
25 // This test doesn't fail - it's informational
26 // The behavior depends on whether Ollama is actually running
27}
28
29func TestGetModels(t *testing.T) {
30 if !IsInstalled() {
31 t.Skip("Ollama is not installed, skipping GetModels test")
32 }
33
34 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
35 defer cancel()
36
37 // Ensure Ollama is running
38 if !IsRunning(ctx) {
39 t.Log("Ollama is not running, attempting to start...")
40 if err := StartOllamaService(ctx); err != nil {
41 t.Fatalf("Failed to start Ollama service: %v", err)
42 }
43 defer cleanupProcesses()
44 }
45
46 models, err := GetModels(ctx)
47 if err != nil {
48 t.Fatalf("Failed to get models: %v", err)
49 }
50
51 t.Logf("✓ Found %d models:", len(models))
52 for _, model := range models {
53 t.Logf(" - %s (context: %d, max_tokens: %d)",
54 model.ID, model.ContextWindow, model.DefaultMaxTokens)
55 }
56}
57
58func TestGetRunningModels(t *testing.T) {
59 if !IsInstalled() {
60 t.Skip("Ollama is not installed, skipping GetRunningModels test")
61 }
62
63 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
64 defer cancel()
65
66 // Ensure Ollama is running
67 if !IsRunning(ctx) {
68 t.Log("Ollama is not running, attempting to start...")
69 if err := StartOllamaService(ctx); err != nil {
70 t.Fatalf("Failed to start Ollama service: %v", err)
71 }
72 defer cleanupProcesses()
73 }
74
75 runningModels, err := GetRunningModels(ctx)
76 if err != nil {
77 t.Fatalf("Failed to get running models: %v", err)
78 }
79
80 t.Logf("✓ Found %d running models:", len(runningModels))
81 for _, model := range runningModels {
82 t.Logf(" - %s (size: %d bytes)", model.Name, model.Size)
83 }
84}
85
86func TestIsModelLoaded(t *testing.T) {
87 if !IsInstalled() {
88 t.Skip("Ollama is not installed, skipping IsModelLoaded test")
89 }
90
91 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
92 defer cancel()
93
94 // Ensure Ollama is running
95 if !IsRunning(ctx) {
96 t.Log("Ollama is not running, attempting to start...")
97 if err := StartOllamaService(ctx); err != nil {
98 t.Fatalf("Failed to start Ollama service: %v", err)
99 }
100 defer cleanupProcesses()
101 }
102
103 // Get available models first
104 models, err := GetModels(ctx)
105 if err != nil {
106 t.Fatalf("Failed to get models: %v", err)
107 }
108
109 if len(models) == 0 {
110 t.Skip("No models available, skipping IsModelLoaded test")
111 }
112
113 testModel := models[0].ID
114 t.Logf("Testing model: %s", testModel)
115
116 loaded, err := IsModelLoaded(ctx, testModel)
117 if err != nil {
118 t.Fatalf("Failed to check if model is loaded: %v", err)
119 }
120
121 if loaded {
122 t.Logf("✓ Model %s is loaded", testModel)
123 } else {
124 t.Logf("✗ Model %s is not loaded", testModel)
125 }
126}
127
128func TestGetProvider(t *testing.T) {
129 if !IsInstalled() {
130 t.Skip("Ollama is not installed, skipping GetProvider test")
131 }
132
133 ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
134 defer cancel()
135
136 // Ensure Ollama is running
137 if !IsRunning(ctx) {
138 t.Log("Ollama is not running, attempting to start...")
139 if err := StartOllamaService(ctx); err != nil {
140 t.Fatalf("Failed to start Ollama service: %v", err)
141 }
142 defer cleanupProcesses()
143 }
144
145 provider, err := GetProvider(ctx)
146 if err != nil {
147 t.Fatalf("Failed to get provider: %v", err)
148 }
149
150 if provider.Name != "Ollama" {
151 t.Errorf("Expected provider name to be 'Ollama', got '%s'", provider.Name)
152 }
153
154 if provider.ID != "ollama" {
155 t.Errorf("Expected provider ID to be 'ollama', got '%s'", provider.ID)
156 }
157
158 t.Logf("✓ Provider: %s (ID: %s) with %d models",
159 provider.Name, provider.ID, len(provider.Models))
160}
161
162func TestGetContextWindow(t *testing.T) {
163 tests := []struct {
164 family string
165 expected int64
166 }{
167 {"llama", 131072},
168 {"mistral", 32768},
169 {"gemma", 8192},
170 {"qwen", 131072},
171 {"qwen2", 131072},
172 {"phi", 131072},
173 {"codellama", 16384},
174 {"unknown", 8192},
175 }
176
177 for _, tt := range tests {
178 t.Run(tt.family, func(t *testing.T) {
179 result := getContextWindow(tt.family)
180 if result != tt.expected {
181 t.Errorf("getContextWindow(%s) = %d, expected %d",
182 tt.family, result, tt.expected)
183 }
184 })
185 }
186}
187
188func TestSupportsImages(t *testing.T) {
189 tests := []struct {
190 family string
191 expected bool
192 }{
193 {"llama-vision", true},
194 {"llava", true},
195 {"llama", false},
196 {"mistral", false},
197 {"unknown", false},
198 }
199
200 for _, tt := range tests {
201 t.Run(tt.family, func(t *testing.T) {
202 result := supportsImages(tt.family)
203 if result != tt.expected {
204 t.Errorf("supportsImages(%s) = %v, expected %v",
205 tt.family, result, tt.expected)
206 }
207 })
208 }
209}
210
211// Benchmark tests for client functions
212func BenchmarkIsRunning(b *testing.B) {
213 if !IsInstalled() {
214 b.Skip("Ollama is not installed")
215 }
216
217 ctx := context.Background()
218
219 for i := 0; i < b.N; i++ {
220 IsRunning(ctx)
221 }
222}
223
224func BenchmarkGetModels(b *testing.B) {
225 if !IsInstalled() {
226 b.Skip("Ollama is not installed")
227 }
228
229 ctx := context.Background()
230
231 // Ensure Ollama is running for benchmark
232 if !IsRunning(ctx) {
233 b.Skip("Ollama is not running")
234 }
235
236 for i := 0; i < b.N; i++ {
237 GetModels(ctx)
238 }
239}