llmgw/internal/cli/bench_test.go

package cli

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/adminapi"
)

// benchCluster is a gateway taking the client's key, a llama.cpp replica
// whose /metrics takes only its engine key, and an admin API that lists the
// replica with that key, as the control plane does.
type benchCluster struct {
	gateway, engine, admin *httptest.Server
	metricsAuth            atomic.Value
}

func newBenchCluster(t *testing.T) *benchCluster {
	t.Helper()
	c := &benchCluster{}
	c.metricsAuth.Store("")
	c.gateway = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer client" {
			http.Error(w, "no", http.StatusUnauthorized)
			return
		}
		w.Header().Set("Content-Type", "text/event-stream")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n")
		fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n")
		fmt.Fprint(w, "data: [DONE]\n\n")
	}))
	c.engine = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		c.metricsAuth.Store(r.Header.Get("Authorization"))
		if r.URL.Path != "/metrics" || r.Header.Get("Authorization") != "Bearer engine-key" {
			http.Error(w, `{"error":"Invalid API Key"}`, http.StatusUnauthorized)
			return
		}
		fmt.Fprint(w, "llamacpp:requests_processing 1\nllamacpp:requests_deferred 3\n")
	}))
	c.admin = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_ = json.NewEncoder(w).Encode([]adminapi.ReplicaView{
			{ID: "r1", Model: "m", Profile: "cpu", Endpoint: c.engine.URL, EngineKey: "engine-key"},
		})
	}))
	t.Cleanup(func() { c.gateway.Close(); c.engine.Close(); c.admin.Close() })
	return c
}

func (c *benchCluster) bench(t *testing.T, args ...string) (string, error) {
	t.Helper()
	t.Setenv("LLMGW_ADMIN_TOKEN", "t")
	asJSON := false
	cmd := newBenchCmd(func() *ctlClient { return newCtlClient(c.admin.URL) }, &asJSON)
	var out bytes.Buffer
	cmd.SetOut(&out)
	cmd.SetErr(&out)
	cmd.SetArgs(append([]string{"m", "--concurrency", "1", "--duration", "1500ms", "--ttft-p99", "1m", "--tpot-p99", "1m"}, args...))
	err := cmd.ExecuteContext(context.Background())
	return out.String(), err
}

func TestBenchReadsAGuardedEnginesMetricsWithTheKeyTheAdminAPIGives(t *testing.T) {
	// Through the gateway with a client's key, and the replica's metrics on
	// the side: the client's key is not the engine's, and the engine's is
	// found by the replica at --metrics' address, not at the target's.
	c := newBenchCluster(t)
	out, err := c.bench(t, "--target", c.gateway.URL, "--key", "client",
		"--metrics", c.engine.URL+"/metrics", "--engine", "llamacpp")
	if err != nil {
		t.Fatalf("%v\n%s", err, out)
	}
	if got := c.metricsAuth.Load(); got != "Bearer engine-key" {
		t.Errorf("metrics read with %q", got)
	}
	if !strings.Contains(out, "target_waiting: 3") {
		t.Errorf("the queue depth read was not suggested:\n%s", out)
	}
}

func TestBenchSaysWhyAnEnginesMetricsCannotBeRead(t *testing.T) {
	// A key the engine refuses fails the run before any load, saying where
	// the right one comes from, rather than reporting an idle engine.
	c := newBenchCluster(t)
	out, err := c.bench(t, "--target", c.gateway.URL, "--key", "client",
		"--metrics", c.engine.URL+"/metrics", "--engine", "llamacpp", "--metrics-key", "wrong")
	if err == nil || !strings.Contains(err.Error(), "refused the key") || !strings.Contains(err.Error(), "--metrics-key") {
		t.Fatalf("expected the refusal explained, got %v\n%s", err, out)
	}
}