1package main
2
3import (
4 "context"
5 "fmt"
6 "os"
7 "time"
8
9 "github.com/charmbracelet/crush/internal/ollama"
10)
11
12func main() {
13 fmt.Println("๐งช Ollama Test Suite")
14 fmt.Println("===================")
15
16 // Test 1: Check if Ollama is installed
17 fmt.Print("1. Checking if Ollama is installed... ")
18 if ollama.IsInstalled() {
19 fmt.Println("โ
PASS")
20 } else {
21 fmt.Println("โ FAIL - Ollama is not installed")
22 fmt.Println(" Please install Ollama from https://ollama.com")
23 os.Exit(1)
24 }
25
26 // Test 2: Check if Ollama is running
27 fmt.Print("2. Checking if Ollama is running... ")
28 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
29 defer cancel()
30
31 if ollama.IsRunning(ctx) {
32 fmt.Println("โ
PASS")
33 } else {
34 fmt.Println("โ FAIL - Ollama is not running")
35
36 // Test 3: Try to start Ollama service
37 fmt.Print("3. Attempting to start Ollama service... ")
38 ctx2, cancel2 := context.WithTimeout(context.Background(), 15*time.Second)
39 defer cancel2()
40
41 if err := ollama.StartOllamaService(ctx2); err != nil {
42 fmt.Printf("โ FAIL - %v\n", err)
43 os.Exit(1)
44 }
45 fmt.Println("โ
PASS")
46
47 // Verify it's now running
48 fmt.Print("4. Verifying Ollama is now running... ")
49 if ollama.IsRunning(ctx2) {
50 fmt.Println("โ
PASS")
51 } else {
52 fmt.Println("โ FAIL - Service started but not responding")
53 os.Exit(1)
54 }
55 }
56
57 // Test 4: Get available models
58 fmt.Print("5. Getting available models... ")
59 ctx3, cancel3 := context.WithTimeout(context.Background(), 10*time.Second)
60 defer cancel3()
61
62 models, err := ollama.GetModels(ctx3)
63 if err != nil {
64 fmt.Printf("โ FAIL - %v\n", err)
65 os.Exit(1)
66 }
67 fmt.Printf("โ
PASS (%d models found)\n", len(models))
68
69 // Display models
70 if len(models) > 0 {
71 fmt.Println("\n๐ Available Models:")
72 for i, model := range models {
73 fmt.Printf(" %d. %s\n", i+1, model.ID)
74 fmt.Printf(" Context: %d tokens, Max: %d tokens\n",
75 model.ContextWindow, model.DefaultMaxTokens)
76 }
77 } else {
78 fmt.Println("\nโ ๏ธ No models found. You may need to download some models first.")
79 fmt.Println(" Example: ollama pull llama3.2:3b")
80 }
81
82 // Test 5: Get provider
83 fmt.Print("\n6. Getting Ollama provider... ")
84 provider, err := ollama.GetProvider(ctx3)
85 if err != nil {
86 fmt.Printf("โ FAIL - %v\n", err)
87 os.Exit(1)
88 }
89 fmt.Printf("โ
PASS (%s with %d models)\n", provider.Name, len(provider.Models))
90
91 // Test 6: Test model loading check
92 if len(models) > 0 {
93 testModel := models[0].ID
94 fmt.Printf("7. Checking if model '%s' is loaded... ", testModel)
95
96 loaded, err := ollama.IsModelLoaded(ctx3, testModel)
97 if err != nil {
98 fmt.Printf("โ FAIL - %v\n", err)
99 } else if loaded {
100 fmt.Println("โ
PASS (model is loaded)")
101 } else {
102 fmt.Println("โ ๏ธ PASS (model is not loaded)")
103 }
104 }
105
106 fmt.Println("\n๐ All tests completed successfully!")
107 fmt.Println("\nTo run individual tests:")
108 fmt.Println(" go test ./internal/ollama -v")
109 fmt.Println("\nTo run benchmarks:")
110 fmt.Println(" go test ./internal/ollama -bench=.")
111}