llmgw/internal/agent/supervisor/health_test.go

package supervisor

import (
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// engineStub serves configurable /health and /v1/models responses.
func engineStub(t *testing.T, healthy bool, models ...string) *httptest.Server {
	t.Helper()
	mux := http.NewServeMux()
	mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
		if !healthy {
			http.Error(w, "loading", http.StatusServiceUnavailable)
			return
		}
		w.WriteHeader(http.StatusOK)
	})
	mux.HandleFunc("GET /v1/models", func(w http.ResponseWriter, _ *http.Request) {
		var b strings.Builder
		b.WriteString(`{"object":"list","data":[`)
		for i, m := range models {
			if i > 0 {
				b.WriteString(",")
			}
			b.WriteString(`{"id":"` + m + `"}`)
		}
		b.WriteString(`]}`)
		io.WriteString(w, b.String())
	})
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	return srv
}

func TestReadyWhenServingModel(t *testing.T) {
	srv := engineStub(t, true, "llama-3.3-70b")
	if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b", ""); err != nil {
		t.Errorf("a serving engine was reported unready: %v", err)
	}
}

func TestNotReadyWhileLoading(t *testing.T) {
	// The case that matters: the process is up but the weights are not loaded.
	// Routing on /health alone would send traffic into a loading model.
	srv := engineStub(t, false, "llama-3.3-70b")
	err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b", "")
	if err == nil {
		t.Fatal("an unhealthy engine was reported ready")
	}
	if !strings.Contains(err.Error(), "health") {
		t.Errorf("err=%v", err)
	}
}

func TestNotReadyWhenModelMissing(t *testing.T) {
	// /health passes but the model is not listed: the engine is up and serving
	// something else, or still registering.
	srv := engineStub(t, true, "some-other-model")
	err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "llama-3.3-70b", "")
	if err == nil {
		t.Fatal("a model that is not served was reported ready")
	}
	if !strings.Contains(err.Error(), "does not serve") {
		t.Errorf("err=%v", err)
	}
}

func TestModelCheckCanBeSkipped(t *testing.T) {
	srv := engineStub(t, true)
	h := NewHTTPHealth(time.Second, WithoutModelCheck())
	if err := h.Ready(t.Context(), srv.URL, "anything", ""); err != nil {
		t.Errorf("the model check was not skipped: %v", err)
	}
}

func TestUnreachableEngine(t *testing.T) {
	h := NewHTTPHealth(200 * time.Millisecond)
	if err := h.Ready(t.Context(), "http://127.0.0.1:1", "m", ""); err == nil {
		t.Error("an unreachable engine was reported ready")
	}
}

func TestTheHealthCheckSendsTheEnginesKey(t *testing.T) {
	// An engine started with a key refuses /v1/models without it, and a
	// shard that is serving would read as never ready.
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") != "Bearer k1" {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		_, _ = w.Write([]byte(`{"data":[{"id":"m"}]}`))
	}))
	defer srv.Close()
	if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "m", "k1"); err != nil {
		t.Errorf("with the key: %v", err)
	}
	if err := NewHTTPHealth(time.Second).Ready(t.Context(), srv.URL, "m", ""); err == nil {
		t.Error("without the key the engine refused, and that read as ready")
	}
}

func TestNotReadyWhileTheHealthBodySaysLoading(t *testing.T) {
	// FreeToken answers 200 from the moment its HTTP server is up and lists
	// its model meanwhile; only the body says the weights are still loading.
	body := `{"status":"loading","phase":"weights","progress":{"done_bytes":1,"total_bytes":9},"model":"glm"}`
	mux := http.NewServeMux()
	mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) { io.WriteString(w, body) })
	mux.HandleFunc("GET /v1/models", func(w http.ResponseWriter, _ *http.Request) {
		io.WriteString(w, `{"data":[{"id":"glm"}]}`)
	})
	srv := httptest.NewServer(mux)
	t.Cleanup(srv.Close)
	h := NewHTTPHealth(time.Second)
	if err := h.Ready(t.Context(), srv.URL, "glm", ""); err == nil || !strings.Contains(err.Error(), "loading") {
		t.Errorf("a loading engine was reported ready, or not said to be loading: %v", err)
	}
	body = `{"status":"ok","model":"glm","uptime_s":3}`
	if err := h.Ready(t.Context(), srv.URL, "glm", ""); err != nil {
		t.Errorf("a serving engine was reported unready: %v", err)
	}
	// Text, as some engines answer: not a status, and not a failure.
	body = "OK"
	if err := h.Ready(t.Context(), srv.URL, "glm", ""); err != nil {
		t.Errorf("a plain-text health answer was taken for a failure: %v", err)
	}
}