1package ollama
2
3import (
4 "context"
5 "os/exec"
6 "testing"
7 "time"
8)
9
10func TestProcessManagementWithRealModel(t *testing.T) {
11 if !IsInstalled() {
12 t.Skip("Ollama is not installed, skipping process management test")
13 }
14
15 ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
16 defer cancel()
17
18 // Start with a clean state
19 originallyRunning := IsRunning(ctx)
20 t.Logf("Ollama originally running: %v", originallyRunning)
21
22 // If Ollama wasn't running, we'll start it and be responsible for cleanup
23 var shouldCleanup bool
24 if !originallyRunning {
25 shouldCleanup = true
26 t.Log("Starting Ollama service...")
27
28 if err := StartService(ctx); err != nil {
29 t.Fatalf("Failed to start Ollama service: %v", err)
30 }
31
32 if !IsRunning(ctx) {
33 t.Fatal("Started Ollama service but it's not running")
34 }
35
36 t.Log("✓ Ollama service started successfully")
37 }
38
39 // Get available models
40 models, err := GetModels(ctx)
41 if err != nil {
42 t.Fatalf("Failed to get models: %v", err)
43 }
44
45 if len(models) == 0 {
46 t.Skip("No models available, skipping model loading test")
47 }
48
49 // Choose a test model (prefer smaller models)
50 testModel := models[0].Name
51 for _, model := range models {
52 if model.Name == "phi3:3.8b" || model.Name == "llama3.2:3b" {
53 testModel = model.Name
54 break
55 }
56 }
57
58 t.Logf("Testing with model: %s", testModel)
59
60 // Test 1: Load model
61 t.Log("Loading model...")
62 startTime := time.Now()
63
64 if err := EnsureModelLoaded(ctx, testModel); err != nil {
65 t.Fatalf("Failed to load model: %v", err)
66 }
67
68 loadTime := time.Since(startTime)
69 t.Logf("✓ Model loaded in %v", loadTime)
70
71 // Verify model is running
72 running, err := IsModelRunning(ctx, testModel)
73 if err != nil {
74 t.Fatalf("Failed to check if model is running: %v", err)
75 }
76
77 if !running {
78 t.Fatal("Model should be running but isn't")
79 }
80
81 t.Log("✓ Model is confirmed running")
82
83 // Test 2: Immediate cleanup after loading
84 t.Log("Testing immediate cleanup after model load...")
85
86 cleanupStart := time.Now()
87 cleanup()
88 cleanupTime := time.Since(cleanupStart)
89
90 t.Logf("✓ Cleanup completed in %v", cleanupTime)
91
92 // Give cleanup time to take effect
93 time.Sleep(2 * time.Second)
94
95 // Test 3: Verify cleanup worked
96 if shouldCleanup {
97 // If we started Ollama, it should be stopped
98 if IsRunning(ctx) {
99 t.Error("❌ Ollama service should be stopped after cleanup but it's still running")
100 } else {
101 t.Log("✓ Ollama service properly stopped after cleanup")
102 }
103 } else {
104 // If Ollama was already running, it should still be running but model should be stopped
105 if !IsRunning(ctx) {
106 t.Error("❌ Ollama service should still be running but it's stopped")
107 } else {
108 t.Log("✓ Ollama service still running (as expected)")
109
110 // Check if model is still loaded
111 running, err := IsModelRunning(ctx, testModel)
112 if err != nil {
113 t.Errorf("Failed to check model status after cleanup: %v", err)
114 } else if running {
115 t.Error("❌ Model should be stopped after cleanup but it's still running")
116 } else {
117 t.Log("✓ Model properly stopped after cleanup")
118 }
119 }
120 }
121
122 // Test 4: Test cleanup idempotency
123 t.Log("Testing cleanup idempotency...")
124 cleanup()
125 cleanup()
126 cleanup()
127 t.Log("✓ Multiple cleanup calls handled safely")
128}
129
130func TestCleanupWithMockProcess(t *testing.T) {
131 // Test cleanup mechanism with a mock process that simulates Ollama
132 cmd := exec.Command("sleep", "30")
133 if err := cmd.Start(); err != nil {
134 t.Fatalf("Failed to start mock process: %v", err)
135 }
136
137 pid := cmd.Process.Pid
138 t.Logf("Started mock process with PID: %d", pid)
139
140 // Simulate what happens in our process manager
141 processManager.mu.Lock()
142 processManager.ollamaProcess = cmd
143 processManager.crushStartedOllama = true
144 processManager.mu.Unlock()
145
146 // Test cleanup
147 t.Log("Testing cleanup with mock process...")
148 stopOllamaService()
149
150 // Verify process was terminated
151 if cmd.ProcessState != nil && cmd.ProcessState.Exited() {
152 t.Log("✓ Mock process was successfully terminated")
153 } else {
154 // Process might still be terminating
155 time.Sleep(100 * time.Millisecond)
156 if cmd.ProcessState != nil && cmd.ProcessState.Exited() {
157 t.Log("✓ Mock process was successfully terminated")
158 } else {
159 t.Error("❌ Mock process was not terminated")
160 }
161 }
162}
163
164func TestSetupCleanup(t *testing.T) {
165 // Test that setupCleanup can be called without panicking
166 defer func() {
167 if r := recover(); r != nil {
168 t.Fatalf("setupCleanup panicked: %v", r)
169 }
170 }()
171
172 // This should not panic and should be safe to call multiple times
173 setupCleanup()
174 t.Log("✓ setupCleanup completed without panic")
175}
176
177func TestStopModel(t *testing.T) {
178 if !IsInstalled() {
179 t.Skip("Ollama is not installed, skipping stopModel test")
180 }
181
182 ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
183 defer cancel()
184
185 // Ensure Ollama is running
186 if err := EnsureRunning(ctx); err != nil {
187 t.Fatalf("Failed to ensure Ollama is running: %v", err)
188 }
189
190 // Get available models
191 models, err := GetModels(ctx)
192 if err != nil {
193 t.Fatalf("Failed to get models: %v", err)
194 }
195
196 if len(models) == 0 {
197 t.Skip("No models available, skipping stopModel test")
198 }
199
200 testModel := models[0].Name
201 t.Logf("Testing stop with model: %s", testModel)
202
203 // Load the model first
204 if err := EnsureModelLoaded(ctx, testModel); err != nil {
205 t.Fatalf("Failed to load model: %v", err)
206 }
207
208 // Verify it's running
209 running, err := IsModelRunning(ctx, testModel)
210 if err != nil {
211 t.Fatalf("Failed to check if model is running: %v", err)
212 }
213
214 if !running {
215 t.Fatal("Model should be running but isn't")
216 }
217
218 // Test stopping the model
219 t.Log("Stopping model...")
220 if err := stopModel(ctx, testModel); err != nil {
221 t.Fatalf("Failed to stop model: %v", err)
222 }
223
224 // Give it time to stop
225 time.Sleep(2 * time.Second)
226
227 // Verify it's stopped
228 running, err = IsModelRunning(ctx, testModel)
229 if err != nil {
230 t.Fatalf("Failed to check if model is running after stop: %v", err)
231 }
232
233 if running {
234 t.Error("❌ Model should be stopped but it's still running")
235 } else {
236 t.Log("✓ Model successfully stopped")
237 }
238
239 // Cleanup
240 defer func() {
241 if processManager.crushStartedOllama {
242 cleanup()
243 }
244 }()
245}