First AI generated version

2026-06-11 20:01:48 +02:00
commit c7d14ffee7
26 changed files with 1884 additions and 0 deletions
@@ -0,0 +1,159 @@
+package gollama_test
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	gollama "gitea.neitzel.de/konrad/go-ollama"
+)
+
+func TestPing(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method == http.MethodHead && r.URL.Path == "/" {
+			w.WriteHeader(http.StatusOK)
+			return
+		}
+		http.NotFound(w, r)
+	}))
+	defer srv.Close()
+
+	client, err := gollama.New(srv.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := client.Ping(context.Background()); err != nil {
+		t.Fatalf("ping: %v", err)
+	}
+}
+
+func TestModelNames(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/api/tags" {
+			_ = json.NewEncoder(w).Encode(map[string]any{
+				"models": []map[string]any{
+					{"name": "llama3.2", "model": "llama3.2", "size": 100},
+					{"name": "mistral", "model": "mistral", "size": 200},
+				},
+			})
+			return
+		}
+		http.NotFound(w, r)
+	}))
+	defer srv.Close()
+
+	client, err := gollama.New(srv.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	names, err := client.ModelNames(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(names) != 2 || names[0] != "llama3.2" || names[1] != "mistral" {
+		t.Fatalf("names = %v", names)
+	}
+}
+
+func TestComplete(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/api/generate" {
+			http.NotFound(w, r)
+			return
+		}
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"model":    "test",
+			"response": "hello world",
+			"done":     true,
+		})
+	}))
+	defer srv.Close()
+
+	client, err := gollama.New(srv.URL, gollama.WithDefaultModel("test"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	result, err := client.Complete(context.Background(), "hi")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if result.Text != "hello world" {
+		t.Fatalf("text = %q", result.Text)
+	}
+}
+
+func TestChatMessages(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/api/chat" {
+			http.NotFound(w, r)
+			return
+		}
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"model": "test",
+			"message": map[string]any{
+				"role":    "assistant",
+				"content": "chat reply",
+			},
+			"done": true,
+		})
+	}))
+	defer srv.Close()
+
+	client, err := gollama.New(srv.URL, gollama.WithDefaultModel("test"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	result, err := client.ChatMessages(context.Background(), []gollama.Message{
+		gollama.User("hello"),
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if result.Text != "chat reply" {
+		t.Fatalf("text = %q", result.Text)
+	}
+}
+
+func TestSession(t *testing.T) {
+	calls := 0
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != "/api/chat" {
+			http.NotFound(w, r)
+			return
+		}
+		calls++
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"model": "test",
+			"message": map[string]any{
+				"role":    "assistant",
+				"content": "reply",
+			},
+			"done": true,
+		})
+	}))
+	defer srv.Close()
+
+	client, err := gollama.New(srv.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	session, err := client.NewSession("test", gollama.WithSessionSystem("be brief"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := session.Send(context.Background(), "one"); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := session.Send(context.Background(), "two"); err != nil {
+		t.Fatal(err)
+	}
+	if calls != 2 {
+		t.Fatalf("calls = %d", calls)
+	}
+	msgs := session.Messages()
+	if len(msgs) < 3 {
+		t.Fatalf("messages = %d", len(msgs))
+	}
+}