steve
7cd7eaff8b
Allow configuring how long the worker model stays resident on the Ollama
target after a request via FOREMAN_KEEP_ALIVE env var. Accepts Ollama
duration strings ("-1" forever, "0" unload, "15m", "1h", etc). Defaults
to "-1" (pin forever). The embedder warm-up is unaffected and always
uses keep_alive=-1.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
470 lines
14 KiB
Go
470 lines
14 KiB
Go
package client_test
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"path/filepath"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.stevedudenhoeffer.com/steve/foreman/client"
|
|
"gitea.stevedudenhoeffer.com/steve/foreman/internal/config"
|
|
"gitea.stevedudenhoeffer.com/steve/foreman/internal/ollama"
|
|
"gitea.stevedudenhoeffer.com/steve/foreman/internal/server"
|
|
"gitea.stevedudenhoeffer.com/steve/foreman/internal/store"
|
|
"gitea.stevedudenhoeffer.com/steve/foreman/internal/webhook"
|
|
"gitea.stevedudenhoeffer.com/steve/foreman/internal/worker"
|
|
)
|
|
|
|
// newTestForeman creates a fully wired foreman server (store, worker, dispatcher)
|
|
// backed by a temp SQLite database and stub Ollama client. Returns the httptest
|
|
// server URL and a cleanup function.
|
|
//
|
|
// Why: the client tests exercise the real HTTP API to catch integration bugs.
|
|
// What: wires store, inventory, notifier, dispatcher, worker, and server into an
|
|
// httptest.Server.
|
|
// Test: used as a helper in all client_test.go test functions.
|
|
func newTestForeman(t *testing.T, ollamaClient ollama.Client, webhookSecret string) string {
|
|
t.Helper()
|
|
|
|
dbPath := filepath.Join(t.TempDir(), "test.db")
|
|
st, err := store.Open(dbPath)
|
|
if err != nil {
|
|
t.Fatalf("store.Open: %v", err)
|
|
}
|
|
t.Cleanup(func() { st.Close() })
|
|
|
|
logger := slog.New(slog.NewJSONHandler(io.Discard, nil))
|
|
inv := ollama.NewModelInventory(ollamaClient, logger)
|
|
if err := inv.Refresh(context.Background()); err != nil {
|
|
t.Fatalf("inv.Refresh: %v", err)
|
|
}
|
|
|
|
notifier := worker.NewNotifier()
|
|
dispatcher := webhook.NewDispatcher(webhookSecret, logger)
|
|
w := worker.New(st, ollamaClient, inv, notifier, dispatcher, logger, "-1")
|
|
|
|
cfg := config.Config{
|
|
OllamaURL: "http://localhost:11434",
|
|
MaxAttempts: 3,
|
|
JobTTL: 24 * time.Hour,
|
|
WebhookSecret: webhookSecret,
|
|
}
|
|
|
|
srv := server.New(cfg, st, ollamaClient, inv, notifier, w, dispatcher, logger)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
go w.Run(ctx)
|
|
|
|
ts := httptest.NewServer(srv.Handler())
|
|
t.Cleanup(ts.Close)
|
|
|
|
return ts.URL
|
|
}
|
|
|
|
func TestSubmit_HappyPath_Polling(t *testing.T) {
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{
|
|
Models: []ollama.ModelInfo{{Name: "qwen3:30b"}},
|
|
},
|
|
ps: &ollama.PsResponse{},
|
|
chatFunc: func(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error) {
|
|
return &ollama.ChatResponse{
|
|
Model: req.Model,
|
|
Done: true,
|
|
Message: &ollama.Message{Role: "assistant", Content: "Hello from foreman!"},
|
|
}, nil, nil
|
|
},
|
|
}
|
|
|
|
baseURL := newTestForeman(t, ollamaStub, "")
|
|
c := client.New(baseURL, client.WithPollingMode(), client.WithPollInterval(100*time.Millisecond))
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
result, err := c.Submit(ctx, client.SubmitRequest{
|
|
Model: "qwen3:30b",
|
|
Messages: []json.RawMessage{json.RawMessage(`{"role":"user","content":"hi"}`)},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Submit: %v", err)
|
|
}
|
|
|
|
if result.State != "done" {
|
|
t.Errorf("state = %q, want %q", result.State, "done")
|
|
}
|
|
if result.JobID == "" {
|
|
t.Error("job_id should not be empty")
|
|
}
|
|
if result.Model != "qwen3:30b" {
|
|
t.Errorf("model = %q, want %q", result.Model, "qwen3:30b")
|
|
}
|
|
if len(result.Result) == 0 {
|
|
t.Error("result should not be empty on done")
|
|
}
|
|
if len(result.Artifacts) == 0 {
|
|
t.Error("artifacts should include the completion")
|
|
}
|
|
|
|
// Verify the result contains the expected chat response.
|
|
var chatResp ollama.ChatResponse
|
|
if err := json.Unmarshal(result.Result, &chatResp); err != nil {
|
|
t.Fatalf("unmarshal result: %v", err)
|
|
}
|
|
if chatResp.Message == nil || chatResp.Message.Content != "Hello from foreman!" {
|
|
t.Errorf("chat content = %v, want 'Hello from foreman!'", chatResp.Message)
|
|
}
|
|
}
|
|
|
|
func TestSubmit_HappyPath_Webhook(t *testing.T) {
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{
|
|
Models: []ollama.ModelInfo{{Name: "qwen3:30b"}},
|
|
},
|
|
ps: &ollama.PsResponse{},
|
|
chatFunc: func(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error) {
|
|
return &ollama.ChatResponse{
|
|
Model: req.Model,
|
|
Done: true,
|
|
Message: &ollama.Message{Role: "assistant", Content: "webhook response"},
|
|
}, nil, nil
|
|
},
|
|
}
|
|
|
|
baseURL := newTestForeman(t, ollamaStub, "")
|
|
c := client.New(baseURL) // Webhook mode (default)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
result, err := c.Submit(ctx, client.SubmitRequest{
|
|
Model: "qwen3:30b",
|
|
Messages: []json.RawMessage{json.RawMessage(`{"role":"user","content":"hi"}`)},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Submit: %v", err)
|
|
}
|
|
|
|
if result.State != "done" {
|
|
t.Errorf("state = %q, want %q", result.State, "done")
|
|
}
|
|
if result.JobID == "" {
|
|
t.Error("job_id should not be empty")
|
|
}
|
|
|
|
// Verify the result.
|
|
var chatResp ollama.ChatResponse
|
|
if err := json.Unmarshal(result.Result, &chatResp); err != nil {
|
|
t.Fatalf("unmarshal result: %v", err)
|
|
}
|
|
if chatResp.Message == nil || chatResp.Message.Content != "webhook response" {
|
|
t.Errorf("chat content = %v, want 'webhook response'", chatResp.Message)
|
|
}
|
|
}
|
|
|
|
func TestSubmit_FailedJob(t *testing.T) {
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{
|
|
Models: []ollama.ModelInfo{{Name: "qwen3:30b"}},
|
|
},
|
|
ps: &ollama.PsResponse{},
|
|
chatFunc: func(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error) {
|
|
return nil, nil, &ollama.HTTPError{StatusCode: 500, Body: "internal error"}
|
|
},
|
|
}
|
|
|
|
baseURL := newTestForeman(t, ollamaStub, "")
|
|
c := client.New(baseURL, client.WithPollingMode(), client.WithPollInterval(100*time.Millisecond))
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
result, err := c.Submit(ctx, client.SubmitRequest{
|
|
Model: "qwen3:30b",
|
|
Messages: []json.RawMessage{json.RawMessage(`{"role":"user","content":"hi"}`)},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Submit: %v", err)
|
|
}
|
|
|
|
if result.State != "failed" {
|
|
t.Errorf("state = %q, want %q", result.State, "failed")
|
|
}
|
|
if result.Error == "" {
|
|
t.Error("error should not be empty on failed job")
|
|
}
|
|
}
|
|
|
|
func TestSubmit_ContextTimeout(t *testing.T) {
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{
|
|
Models: []ollama.ModelInfo{{Name: "qwen3:30b"}},
|
|
},
|
|
ps: &ollama.PsResponse{},
|
|
chatFunc: func(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error) {
|
|
// Block until context cancelled.
|
|
<-ctx.Done()
|
|
return nil, nil, ctx.Err()
|
|
},
|
|
}
|
|
|
|
baseURL := newTestForeman(t, ollamaStub, "")
|
|
c := client.New(baseURL, client.WithPollingMode(), client.WithPollInterval(50*time.Millisecond))
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
|
|
defer cancel()
|
|
|
|
_, err := c.Submit(ctx, client.SubmitRequest{
|
|
Model: "qwen3:30b",
|
|
Messages: []json.RawMessage{json.RawMessage(`{"role":"user","content":"hi"}`)},
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected error on context timeout")
|
|
}
|
|
}
|
|
|
|
func TestSubmit_AuthToken(t *testing.T) {
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{
|
|
Models: []ollama.ModelInfo{{Name: "qwen3:30b"}},
|
|
},
|
|
ps: &ollama.PsResponse{},
|
|
chatFunc: func(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error) {
|
|
return &ollama.ChatResponse{
|
|
Model: req.Model,
|
|
Done: true,
|
|
Message: &ollama.Message{Role: "assistant", Content: "ok"},
|
|
}, nil, nil
|
|
},
|
|
}
|
|
|
|
// Create a foreman with auth required.
|
|
dbPath := filepath.Join(t.TempDir(), "test.db")
|
|
st, err := store.Open(dbPath)
|
|
if err != nil {
|
|
t.Fatalf("store.Open: %v", err)
|
|
}
|
|
t.Cleanup(func() { st.Close() })
|
|
|
|
logger := slog.New(slog.NewJSONHandler(io.Discard, nil))
|
|
inv := ollama.NewModelInventory(ollamaStub, logger)
|
|
inv.Refresh(context.Background())
|
|
|
|
notifier := worker.NewNotifier()
|
|
dispatcher := webhook.NewDispatcher("", logger)
|
|
w := worker.New(st, ollamaStub, inv, notifier, dispatcher, logger, "-1")
|
|
|
|
cfg := config.Config{
|
|
OllamaURL: "http://localhost:11434",
|
|
Token: "my-secret-token",
|
|
MaxAttempts: 3,
|
|
JobTTL: 24 * time.Hour,
|
|
}
|
|
|
|
srv := server.New(cfg, st, ollamaStub, inv, notifier, w, dispatcher, logger)
|
|
|
|
wCtx, wCancel := context.WithCancel(context.Background())
|
|
t.Cleanup(wCancel)
|
|
go w.Run(wCtx)
|
|
|
|
ts := httptest.NewServer(srv.Handler())
|
|
t.Cleanup(ts.Close)
|
|
|
|
// Without token: should fail.
|
|
c := client.New(ts.URL, client.WithPollingMode(), client.WithPollInterval(100*time.Millisecond))
|
|
ctx := context.Background()
|
|
_, err = c.Submit(ctx, client.SubmitRequest{
|
|
Model: "qwen3:30b",
|
|
Messages: []json.RawMessage{json.RawMessage(`{"role":"user","content":"hi"}`)},
|
|
})
|
|
if err == nil {
|
|
t.Fatal("expected error when no token provided")
|
|
}
|
|
|
|
// With correct token: should succeed.
|
|
c = client.New(ts.URL, client.WithToken("my-secret-token"), client.WithPollingMode(), client.WithPollInterval(100*time.Millisecond))
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
result, err := c.Submit(ctx, client.SubmitRequest{
|
|
Model: "qwen3:30b",
|
|
Messages: []json.RawMessage{json.RawMessage(`{"role":"user","content":"hi"}`)},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Submit with correct token: %v", err)
|
|
}
|
|
if result.State != "done" {
|
|
t.Errorf("state = %q, want %q", result.State, "done")
|
|
}
|
|
}
|
|
|
|
func TestTags(t *testing.T) {
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{
|
|
Models: []ollama.ModelInfo{
|
|
{Name: "qwen3:30b", Size: 19000000000},
|
|
{Name: "nomic-embed-text", Size: 300000000},
|
|
},
|
|
},
|
|
ps: &ollama.PsResponse{},
|
|
}
|
|
|
|
baseURL := newTestForeman(t, ollamaStub, "")
|
|
c := client.New(baseURL)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
models, err := c.Tags(ctx)
|
|
if err != nil {
|
|
t.Fatalf("Tags: %v", err)
|
|
}
|
|
|
|
if len(models) != 2 {
|
|
t.Fatalf("got %d models, want 2", len(models))
|
|
}
|
|
if models[0].Name != "qwen3:30b" {
|
|
t.Errorf("first model = %q, want %q", models[0].Name, "qwen3:30b")
|
|
}
|
|
}
|
|
|
|
func TestEmbed(t *testing.T) {
|
|
embedResp := ollama.EmbedResponse{
|
|
Model: "nomic-embed-text",
|
|
Embeddings: [][]float64{{0.1, 0.2, 0.3}},
|
|
}
|
|
respBytes, _ := json.Marshal(embedResp)
|
|
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{},
|
|
ps: &ollama.PsResponse{},
|
|
rawEmbedFunc: func(ctx context.Context, body []byte) (*http.Response, error) {
|
|
return &http.Response{
|
|
StatusCode: 200,
|
|
Header: http.Header{"Content-Type": {"application/json"}},
|
|
Body: io.NopCloser(bytes.NewReader(respBytes)),
|
|
}, nil
|
|
},
|
|
}
|
|
|
|
baseURL := newTestForeman(t, ollamaStub, "")
|
|
c := client.New(baseURL)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
|
|
resp, err := c.Embed(ctx, client.EmbedRequest{
|
|
Model: "nomic-embed-text",
|
|
Input: json.RawMessage(`"test input"`),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Embed: %v", err)
|
|
}
|
|
|
|
if resp.Model != "nomic-embed-text" {
|
|
t.Errorf("model = %q, want %q", resp.Model, "nomic-embed-text")
|
|
}
|
|
if len(resp.Embeddings) != 1 || len(resp.Embeddings[0]) != 3 {
|
|
t.Errorf("embeddings shape wrong: %v", resp.Embeddings)
|
|
}
|
|
}
|
|
|
|
func TestSubmit_MissingModel(t *testing.T) {
|
|
c := client.New("http://localhost:9999")
|
|
_, err := c.Submit(context.Background(), client.SubmitRequest{})
|
|
if err == nil {
|
|
t.Fatal("expected error for missing model")
|
|
}
|
|
}
|
|
|
|
func TestSubmit_WebhookWithHMAC(t *testing.T) {
|
|
secret := "test-hmac-secret"
|
|
|
|
ollamaStub := &stubOllamaClient{
|
|
tags: &ollama.TagsResponse{
|
|
Models: []ollama.ModelInfo{{Name: "qwen3:30b"}},
|
|
},
|
|
ps: &ollama.PsResponse{},
|
|
chatFunc: func(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error) {
|
|
return &ollama.ChatResponse{
|
|
Model: req.Model,
|
|
Done: true,
|
|
Message: &ollama.Message{Role: "assistant", Content: "hmac verified"},
|
|
}, nil, nil
|
|
},
|
|
}
|
|
|
|
baseURL := newTestForeman(t, ollamaStub, secret)
|
|
c := client.New(baseURL, client.WithWebhookSecret(secret))
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
result, err := c.Submit(ctx, client.SubmitRequest{
|
|
Model: "qwen3:30b",
|
|
Messages: []json.RawMessage{json.RawMessage(`{"role":"user","content":"hi"}`)},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Submit with HMAC: %v", err)
|
|
}
|
|
|
|
if result.State != "done" {
|
|
t.Errorf("state = %q, want %q", result.State, "done")
|
|
}
|
|
}
|
|
|
|
// --- Stub Ollama client for integration tests ---
|
|
|
|
type stubOllamaClient struct {
|
|
tags *ollama.TagsResponse
|
|
tagsErr error
|
|
ps *ollama.PsResponse
|
|
psErr error
|
|
|
|
chatFunc func(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error)
|
|
rawEmbedFunc func(ctx context.Context, body []byte) (*http.Response, error)
|
|
}
|
|
|
|
func (s *stubOllamaClient) Chat(ctx context.Context, req ollama.ChatRequest, stream bool) (*ollama.ChatResponse, <-chan ollama.ChatResponse, error) {
|
|
if s.chatFunc != nil {
|
|
return s.chatFunc(ctx, req, stream)
|
|
}
|
|
return nil, nil, fmt.Errorf("stubOllamaClient.Chat not implemented")
|
|
}
|
|
|
|
func (s *stubOllamaClient) Embed(ctx context.Context, req ollama.EmbedRequest) (*ollama.EmbedResponse, error) {
|
|
return nil, fmt.Errorf("stubOllamaClient.Embed not implemented")
|
|
}
|
|
|
|
func (s *stubOllamaClient) Tags(ctx context.Context) (*ollama.TagsResponse, error) {
|
|
if s.tagsErr != nil {
|
|
return nil, s.tagsErr
|
|
}
|
|
return s.tags, nil
|
|
}
|
|
|
|
func (s *stubOllamaClient) Ps(ctx context.Context) (*ollama.PsResponse, error) {
|
|
if s.psErr != nil {
|
|
return nil, s.psErr
|
|
}
|
|
return s.ps, nil
|
|
}
|
|
|
|
func (s *stubOllamaClient) RawChat(ctx context.Context, body []byte) (*http.Response, error) {
|
|
return nil, fmt.Errorf("stubOllamaClient.RawChat not implemented")
|
|
}
|
|
|
|
func (s *stubOllamaClient) RawEmbed(ctx context.Context, body []byte) (*http.Response, error) {
|
|
if s.rawEmbedFunc != nil {
|
|
return s.rawEmbedFunc(ctx, body)
|
|
}
|
|
return nil, fmt.Errorf("stubOllamaClient.RawEmbed not implemented")
|
|
}
|