86 lines
2.0 KiB
Go
86 lines
2.0 KiB
Go
//go:build integration
|
|
|
|
package gollama_test
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"testing"
|
|
|
|
gollama "gitea.neitzel.de/konrad/go-ollama"
|
|
"gitea.neitzel.de/konrad/go-ollama/internal/integrationtest"
|
|
)
|
|
|
|
func TestIntegrationPing(t *testing.T) {
|
|
client := integrationtest.RequireOllama(t)
|
|
cfg := integrationtest.ConfigForTest(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout)
|
|
defer cancel()
|
|
if err := client.Ping(ctx); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
|
|
func TestIntegrationModelNames(t *testing.T) {
|
|
client := integrationtest.RequireOllama(t)
|
|
cfg := integrationtest.ConfigForTest(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout)
|
|
defer cancel()
|
|
names, err := client.ModelNames(ctx)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
t.Logf("models: %v", names)
|
|
}
|
|
|
|
func TestIntegrationComplete(t *testing.T) {
|
|
client := integrationtest.RequireOllama(t)
|
|
cfg := integrationtest.ConfigForTest(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout)
|
|
defer cancel()
|
|
|
|
model := cfg.Model
|
|
names, err := client.ModelNames(ctx)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(names) == 0 {
|
|
t.Skip("no models available on server")
|
|
}
|
|
if !contains(names, model) {
|
|
model = names[0]
|
|
t.Logf("configured model %q not found, using %q", cfg.Model, model)
|
|
}
|
|
|
|
result, err := client.Complete(ctx, "Reply with exactly one word: ok", gollama.WithModel(model))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if result.Text == "" {
|
|
t.Fatal("empty response")
|
|
}
|
|
t.Logf("response: %q", result.Text)
|
|
}
|
|
|
|
func contains(ss []string, s string) bool {
|
|
for _, v := range ss {
|
|
if v == s {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func TestIntegrationPull(t *testing.T) {
|
|
if os.Getenv("GOLLAMA_TEST_PULL") != "1" {
|
|
t.Skip("set GOLLAMA_TEST_PULL=1 to run pull integration test")
|
|
}
|
|
client := integrationtest.RequireOllama(t)
|
|
cfg := integrationtest.ConfigForTest(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout*4)
|
|
defer cancel()
|
|
if err := client.Pull(ctx, cfg.Model); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|