glpi-neural-brain/internal/ollama/client_test.go

package ollama

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"sync/atomic"
	"testing"
	"time"
)

func poolTestServer(t *testing.T, calls *atomic.Int64, fail bool) *httptest.Server {
	t.Helper()
	return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		switch r.URL.Path {
		case "/api/tags":
			_ = json.NewEncoder(w).Encode(map[string]any{"models": []map[string]any{{"name": "qwen3:8b", "digest": "chat"}, {"name": "embeddinggemma:latest", "digest": "embed"}}})
		case "/api/chat":
			calls.Add(1)
			if fail {
				http.Error(w, "busy", http.StatusServiceUnavailable)
				return
			}
			_ = json.NewEncoder(w).Encode(map[string]any{"message": map[string]any{"content": `{"ok":true}`}})
		case "/api/embed":
			calls.Add(1)
			_ = json.NewEncoder(w).Encode(map[string]any{"embeddings": [][]float64{{0.1, 0.2}}})
		default:
			http.NotFound(w, r)
		}
	}))
}

func TestPoolLeastInflightBalancesSerialRequests(t *testing.T) {
	var aCalls, bCalls atomic.Int64
	a := poolTestServer(t, &aCalls, false)
	defer a.Close()
	b := poolTestServer(t, &bCalls, false)
	defer b.Close()
	c := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "a", URL: a.URL}, {Name: "b", URL: b.URL}}, RoutingMode: "least_inflight", NodeMaxInflight: 1, HealthInterval: time.Minute, FailureCooldown: time.Second, RequestTimeout: time.Second, FailoverEnabled: true, RequireSameModelDigest: true, RequireEmbeddingModel: true}, "qwen3:8b", "embeddinggemma")
	if err := c.Ping(context.Background()); err != nil {
		t.Fatal(err)
	}
	for i := 0; i < 4; i++ {
		var out struct {
			OK bool `json:"ok"`
		}
		if err := c.ChatJSON(context.Background(), "system", "user", map[string]any{"type": "object"}, &out); err != nil {
			t.Fatal(err)
		}
	}
	if aCalls.Load() != 2 || bCalls.Load() != 2 {
		t.Fatalf("distribution a=%d b=%d", aCalls.Load(), bCalls.Load())
	}
}

func TestPoolFailsOverOnRetryableError(t *testing.T) {
	var badCalls, goodCalls atomic.Int64
	bad := poolTestServer(t, &badCalls, true)
	defer bad.Close()
	good := poolTestServer(t, &goodCalls, false)
	defer good.Close()
	c := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "a-bad", URL: bad.URL}, {Name: "b-good", URL: good.URL}}, RoutingMode: "least_inflight", NodeMaxInflight: 1, HealthInterval: time.Minute, FailureCooldown: time.Second, RequestTimeout: time.Second, FailoverEnabled: true, FailoverAttempts: 2, RequireSameModelDigest: true, RequireEmbeddingModel: true}, "qwen3:8b", "embeddinggemma")
	if err := c.Ping(context.Background()); err != nil {
		t.Fatal(err)
	}
	var out struct {
		OK bool `json:"ok"`
	}
	if err := c.ChatJSON(context.Background(), "system", "user", map[string]any{"type": "object"}, &out); err != nil {
		t.Fatal(err)
	}
	if badCalls.Load() != 1 || goodCalls.Load() != 1 {
		t.Fatalf("failover bad=%d good=%d", badCalls.Load(), goodCalls.Load())
	}
}

func TestLowPriorityWaiterYieldsToNormalWaiter(t *testing.T) {
	client := NewPool(PoolConfig{Nodes: []NodeConfig{{Name: "only", URL: "http://unused"}}, RoutingMode: "least_inflight", NodeMaxInflight: 1}, "qwen3:8b", "embeddinggemma")
	client.mu.Lock()
	node := client.nodes[0]
	node.healthy = true
	node.compatible = true
	node.chatModel = true
	node.embeddingModel = true
	node.inflight = 1 // hold capacity until both waiters are registered
	client.healthReady = true
	client.mu.Unlock()

	type result struct {
		name string
		node *nodeState
		err  error
	}
	results := make(chan result, 2)
	lowCtx, lowCancel := context.WithTimeout(WithLowPriority(context.Background()), 2*time.Second)
	defer lowCancel()
	go func() {
		n, err := client.acquireNode(lowCtx, "chat", map[*nodeState]bool{})
		results <- result{name: "low", node: n, err: err}
	}()
	waitForPoolWaiters(t, client, 0, 1)

	normalCtx, normalCancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer normalCancel()
	go func() {
		n, err := client.acquireNode(normalCtx, "chat", map[*nodeState]bool{})
		results <- result{name: "normal", node: n, err: err}
	}()
	waitForPoolWaiters(t, client, 1, 1)

	client.mu.Lock()
	node.inflight = 0
	client.mu.Unlock()

	first := <-results
	if first.err != nil {
		t.Fatal(first.err)
	}
	if first.name != "normal" {
		t.Fatalf("low-priority acquisition overtook a normal waiter: first=%s", first.name)
	}
	client.releaseNode(first.node, time.Millisecond, nil)

	second := <-results
	if second.err != nil {
		t.Fatal(second.err)
	}
	if second.name != "low" {
		t.Fatalf("expected low-priority waiter second, got %s", second.name)
	}
	client.releaseNode(second.node, time.Millisecond, nil)
}

func waitForPoolWaiters(t *testing.T, client *Client, normal, low int) {
	t.Helper()
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		client.mu.Lock()
		gotNormal, gotLow := client.normalWaiters, client.lowWaiters
		client.mu.Unlock()
		if gotNormal == normal && gotLow == low {
			return
		}
		time.Sleep(5 * time.Millisecond)
	}
	client.mu.Lock()
	gotNormal, gotLow := client.normalWaiters, client.lowWaiters
	client.mu.Unlock()
	t.Fatalf("waiter counters did not reach normal=%d low=%d; got normal=%d low=%d", normal, low, gotNormal, gotLow)
}