test(distributed): add admin session helper to the cluster harness

The register handler answers 201 both for "user created, here is your
session" and for "this email already exists", so the status code cannot
tell a fresh registration from a repeat one. Key on the session cookie
instead and fall through to login when it is absent.

Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

tests/e2e/distributed/cluster/admin.go

@@ -0,0 +1,168 @@
+package cluster
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/cookiejar"
+	"net/url"
+	"time"
+
+	"github.com/mudler/LocalAI/pkg/httpclient"
+)
+
+const (
+	// adminPassword must satisfy core/http/auth's policy (>= 12 chars and a
+	// zxcvbn score of 3 against hints that include "admin" and "localai").
+	// This one scores 3 today; acknowledgeWeakPassword is sent alongside it so
+	// a future tightening of the policy cannot silently break every failover
+	// spec at setup time.
+	adminPassword = "e2e-admin-password"
+	// sessionCookieName mirrors the unexported constant in core/http/auth.
+	// The register handler returns 201 both for "user created, here is your
+	// session" and for "this email already exists" (a deliberate account
+	// enumeration defence), so the status code alone cannot tell the two
+	// apart: the presence of this cookie is the only reliable signal.
+	sessionCookieName = "session"
+	// authRequestTimeout bounds one register/login round trip.
+	authRequestTimeout = 30 * time.Second
+	// bodyExcerptLimit caps how much of an error response is quoted back.
+	bodyExcerptLimit = 512
+)
+
+// ForTestingEmpty returns a Cluster with no processes. It exists so the package's
+// own argument-validation specs do not need to start anything.
+func ForTestingEmpty() *Cluster {
+	return &Cluster{}
+}
+
+// AdminSession registers the admin user on frontend i and returns a client
+// carrying the resulting session cookie. The email matches LOCALAI_ADMIN_EMAIL,
+// which core/http/auth exempts from the approval gate and assigns the admin
+// role, so registration alone yields an active admin session.
+//
+// Call this ONCE per cluster and share the client. Two reasons: the auth
+// endpoints are rate limited to 5 requests per minute per client IP, and every
+// e2e request arrives from 127.0.0.1; and the returned client is already good
+// for every frontend, because sessions live in the shared Postgres auth DB and
+// Go's cookie jar keys cookies by host without the port.
+func (c *Cluster) AdminSession(i int) (*http.Client, error) {
+	base, err := c.frontendBaseURL(i)
+	if err != nil {
+		return nil, err
+	}
+
+	jar, err := cookiejar.New(nil)
+	if err != nil {
+		return nil, fmt.Errorf("creating cookie jar: %w", err)
+	}
+	// httpclient hardens the transport and refuses redirects; the jar is the one
+	// thing it does not configure, and a session cookie is the whole point here.
+	client := httpclient.NewWithTimeout(authRequestTimeout)
+	client.Jar = jar
+
+	credentials := map[string]any{
+		"email":    c.opts.AdminEmail,
+		"password": adminPassword,
+	}
+	registration := map[string]any{
+		"email":                     c.opts.AdminEmail,
+		"password":                  adminPassword,
+		"name":                      "E2E Admin",
+		"acknowledge_weak_password": true,
+	}
+
+	registerStatus, registerBody, err := postJSON(client, base+"/api/auth/register", registration)
+	if err != nil {
+		return nil, fmt.Errorf("registering admin on frontend %d: %w", i, err)
+	}
+	if hasSessionCookie(jar, base) {
+		return client, nil
+	}
+
+	// No cookie means the user already existed (a repeat call against the same
+	// Postgres), or registration was rejected. Log in; on failure the
+	// registration response is the diagnosis, so carry it into the error.
+	loginStatus, loginBody, err := postJSON(client, base+"/api/auth/login", credentials)
+	if err != nil {
+		return nil, fmt.Errorf("logging in admin on frontend %d: %w", i, err)
+	}
+	if loginStatus != http.StatusOK {
+		return nil, fmt.Errorf(
+			"admin login on frontend %d returned %d (%s); registration had returned %d (%s)",
+			i, loginStatus, loginBody, registerStatus, registerBody)
+	}
+	if !hasSessionCookie(jar, base) {
+		return nil, fmt.Errorf("admin login on frontend %d returned 200 but set no %q cookie: %s", i, sessionCookieName, loginBody)
+	}
+	return client, nil
+}
+
+// GetJSON performs an authenticated GET against a frontend and decodes the body.
+func (c *Cluster) GetJSON(client *http.Client, frontend int, path string, out any) error {
+	base, err := c.frontendBaseURL(frontend)
+	if err != nil {
+		return err
+	}
+	resp, err := client.Get(base + path)
+	if err != nil {
+		return fmt.Errorf("GET %s on frontend %d: %w", path, frontend, err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("GET %s on frontend %d returned %d: %s", path, frontend, resp.StatusCode, excerpt(resp.Body))
+	}
+	if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
+		return fmt.Errorf("decoding %s from frontend %d: %w", path, frontend, err)
+	}
+	return nil
+}
+
+// frontendBaseURL validates the index before FrontendURL indexes the slice: a
+// bare index panic in a helper every failover spec calls is far harder to read
+// than a named error.
+func (c *Cluster) frontendBaseURL(i int) (string, error) {
+	if i < 0 || i >= len(c.frontends) {
+		return "", fmt.Errorf("frontend %d out of range (cluster has %d)", i, len(c.frontends))
+	}
+	return c.FrontendURL(i), nil
+}
+
+// postJSON sends body as JSON and returns the status plus an excerpt of the
+// response, closing the body in every path.
+func postJSON(client *http.Client, endpoint string, body any) (int, string, error) {
+	encoded, err := json.Marshal(body)
+	if err != nil {
+		return 0, "", fmt.Errorf("marshalling request body: %w", err)
+	}
+	resp, err := client.Post(endpoint, "application/json", bytes.NewReader(encoded))
+	if err != nil {
+		return 0, "", err
+	}
+	defer func() { _ = resp.Body.Close() }()
+	return resp.StatusCode, excerpt(resp.Body), nil
+}
+
+// hasSessionCookie reports whether the jar holds a usable session for base.
+func hasSessionCookie(jar *cookiejar.Jar, base string) bool {
+	u, err := url.Parse(base)
+	if err != nil {
+		return false
+	}
+	for _, cookie := range jar.Cookies(u) {
+		if cookie.Name == sessionCookieName && cookie.Value != "" {
+			return true
+		}
+	}
+	return false
+}
+
+func excerpt(r io.Reader) string {
+	data, err := io.ReadAll(io.LimitReader(r, bodyExcerptLimit))
+	if err != nil {
+		return fmt.Sprintf("<unreadable body: %v>", err)
+	}
+	return string(bytes.TrimSpace(data))
+}

tests/e2e/distributed/cluster/cluster_test.go

@@ -3,7 +3,9 @@ package cluster_test
 import (
 	"os"
 	"path/filepath"
+	"time"
 
+	"github.com/mudler/LocalAI/pkg/httpclient"
 	"github.com/mudler/LocalAI/tests/e2e/distributed/cluster"
 
 	. "github.com/onsi/ginkgo/v2"
@@ -35,3 +37,23 @@ var _ = Describe("Cluster options", Label("Distributed"), func() {
 		Expect(err.Error()).To(ContainSubstring("at least one frontend"))
 	})
 })
+
+// The HTTP flow inside AdminSession and GetJSON cannot run here: it needs a
+// built local-ai plus real Postgres and NATS, which arrive with the failover
+// suites. These specs cover the argument validation that would otherwise panic
+// on an out-of-range slice index inside a helper every later spec calls.
+var _ = Describe("Admin session", Label("Distributed"), func() {
+	It("reports a clear error when the frontend index is out of range", func() {
+		c := cluster.ForTestingEmpty()
+		_, err := c.AdminSession(3)
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("frontend 3"))
+	})
+
+	It("reports a clear error when GetJSON names a frontend that does not exist", func() {
+		c := cluster.ForTestingEmpty()
+		err := c.GetJSON(httpclient.NewWithTimeout(time.Second), 1, "/api/nodes", &struct{}{})
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("frontend 1"))
+	})
+})