go-ollama/client_test.go

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))
	}
}