131 lines
2.8 KiB
Go
131 lines
2.8 KiB
Go
package integrationtest
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"path/filepath"
|
|
"runtime"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/pelletier/go-toml/v2"
|
|
|
|
gollama "gitea.neitzel.de/konrad/go-ollama"
|
|
)
|
|
|
|
// Config holds settings for integration tests against a live Ollama server.
|
|
type Config struct {
|
|
Host string
|
|
Model string
|
|
Timeout time.Duration
|
|
}
|
|
|
|
type fileConfig struct {
|
|
Host string `toml:"host"`
|
|
Model string `toml:"model"`
|
|
Timeout string `toml:"timeout"`
|
|
}
|
|
|
|
// LoadConfig merges environment variables, local.toml, and defaults.toml.
|
|
func LoadConfig() (Config, error) {
|
|
cfg := Config{
|
|
Host: "http://debian:11434",
|
|
Model: "llama3.2",
|
|
Timeout: 30 * time.Second,
|
|
}
|
|
|
|
if err := mergeFileConfig(&cfg, filepath.Join(configDir(), "defaults.toml")); err != nil {
|
|
return Config{}, err
|
|
}
|
|
_ = mergeFileConfig(&cfg, filepath.Join(configDir(), "local.toml"))
|
|
|
|
if v := os.Getenv("OLLAMA_HOST"); v != "" {
|
|
cfg.Host = v
|
|
}
|
|
if v := os.Getenv("GOLLAMA_TEST_MODEL"); v != "" {
|
|
cfg.Model = v
|
|
}
|
|
if v := os.Getenv("GOLLAMA_TEST_TIMEOUT"); v != "" {
|
|
d, err := time.ParseDuration(v)
|
|
if err != nil {
|
|
return Config{}, err
|
|
}
|
|
cfg.Timeout = d
|
|
}
|
|
|
|
return cfg, nil
|
|
}
|
|
|
|
func mergeFileConfig(cfg *Config, path string) error {
|
|
b, err := os.ReadFile(path)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
var fc fileConfig
|
|
if err := toml.Unmarshal(b, &fc); err != nil {
|
|
return err
|
|
}
|
|
if fc.Host != "" {
|
|
cfg.Host = fc.Host
|
|
}
|
|
if fc.Model != "" {
|
|
cfg.Model = fc.Model
|
|
}
|
|
if fc.Timeout != "" {
|
|
d, err := time.ParseDuration(fc.Timeout)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
cfg.Timeout = d
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func configDir() string {
|
|
_, file, _, ok := runtime.Caller(0)
|
|
if !ok {
|
|
return "testconfig"
|
|
}
|
|
return filepath.Join(filepath.Dir(file), "..", "..", "testconfig")
|
|
}
|
|
|
|
// NewClient creates a gollama client from integration test configuration.
|
|
func NewClient(t *testing.T) *gollama.Client {
|
|
t.Helper()
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("load integration config: %v", err)
|
|
}
|
|
client, err := gollama.New(cfg.Host, gollama.WithDefaultModel(cfg.Model))
|
|
if err != nil {
|
|
t.Fatalf("create client: %v", err)
|
|
}
|
|
return client
|
|
}
|
|
|
|
// ConfigForTest returns the loaded integration configuration.
|
|
func ConfigForTest(t *testing.T) Config {
|
|
t.Helper()
|
|
cfg, err := LoadConfig()
|
|
if err != nil {
|
|
t.Fatalf("load integration config: %v", err)
|
|
}
|
|
return cfg
|
|
}
|
|
|
|
// RequireOllama skips the test when the configured Ollama server is unreachable.
|
|
func RequireOllama(t *testing.T) *gollama.Client {
|
|
t.Helper()
|
|
client := NewClient(t)
|
|
cfg := ConfigForTest(t)
|
|
ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout)
|
|
defer cancel()
|
|
if err := client.Ping(ctx); err != nil {
|
|
t.Skipf("ollama not reachable at %s: %v", cfg.Host, err)
|
|
}
|
|
return client
|
|
}
|