crossbar/internal/health/health_test.go

package health_test

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

	"git.wntrmute.dev/kyle/crossbar/internal/health"
)

// fake is a llama-server stand-in whose /health can be flipped and whose model list is fixed.
type fake struct {
	srv    *httptest.Server
	down   atomic.Bool
	models []string
	hits   atomic.Int32
}

func newFake(t *testing.T, models ...string) *fake {
	f := &fake{models: models}
	mux := http.NewServeMux()
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		f.hits.Add(1)
		if f.down.Load() {
			http.Error(w, "loading", http.StatusServiceUnavailable)
			return
		}
		_, _ = w.Write([]byte(`{"status":"ok"}`))
	})
	mux.HandleFunc("/v1/models", func(w http.ResponseWriter, r *http.Request) {
		type m struct {
			ID string `json:"id"`
		}
		var data []m
		for _, id := range f.models {
			data = append(data, m{ID: id})
		}
		_ = json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})
	})
	f.srv = httptest.NewServer(mux)
	t.Cleanup(f.srv.Close)
	return f
}

func TestFirstPollMakesHealthy(t *testing.T) {
	a := newFake(t, "zeta", "alpha", "alpha")
	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
	if s, ok := tbl.Get("a"); !ok || s.Healthy || len(s.Loaded) != 0 {
		t.Fatalf("before any poll: %+v %v", s, ok)
	}
	tbl.PollOnce(context.Background())
	s, _ := tbl.Get("a")
	if !s.Healthy || s.Consecutive != 1 || s.LastErr != "" || s.LastOK.IsZero() {
		t.Errorf("after one good poll: %+v", s)
	}
	if len(s.Loaded) != 2 || s.Loaded[0] != "alpha" || s.Loaded[1] != "zeta" {
		t.Errorf("Loaded = %v, want sorted, unique [alpha zeta]", s.Loaded)
	}
}

func TestFailureThenRecoveryNeedsTwoPolls(t *testing.T) {
	a := newFake(t, "m")
	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
	ctx := context.Background()
	tbl.PollOnce(ctx)
	a.down.Store(true)
	tbl.PollOnce(ctx)
	s, _ := tbl.Get("a")
	if s.Healthy || s.Consecutive != 0 || s.LastErr == "" {
		t.Fatalf("after failure: %+v", s)
	}
	if len(s.Loaded) != 1 {
		t.Errorf("Loaded is left as last seen; got %v", s.Loaded)
	}
	a.down.Store(false)
	tbl.PollOnce(ctx)
	if s, _ := tbl.Get("a"); s.Healthy || s.Consecutive != 1 {
		t.Errorf("one good poll after a failure must not be healthy yet: %+v", s)
	}
	tbl.PollOnce(ctx)
	if s, _ := tbl.Get("a"); !s.Healthy || s.Consecutive != 2 || s.LastErr != "" {
		t.Errorf("two good polls: %+v", s)
	}
}

func TestMarkDown(t *testing.T) {
	a := newFake(t, "m")
	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
	tbl.PollOnce(context.Background())
	tbl.MarkDown("a", "connection refused")
	s, _ := tbl.Get("a")
	if s.Healthy || s.Consecutive != 0 || s.LastErr != "marked down: connection refused" {
		t.Errorf("after MarkDown: %+v", s)
	}
	tbl.MarkDown("nobody", "x") // unknown hosts are ignored, not a panic
	tbl.PollOnce(context.Background())
	if s, _ := tbl.Get("a"); s.Healthy {
		t.Errorf("one poll after MarkDown must not be healthy: %+v", s)
	}
}

func TestUnreachableAndUnknown(t *testing.T) {
	tbl := health.New(map[string]string{"a": "http://127.0.0.1:1"}, time.Hour, &http.Client{Timeout: time.Second})
	tbl.PollOnce(context.Background())
	s, ok := tbl.Get("a")
	if !ok || s.Healthy || s.LastErr == "" {
		t.Errorf("unreachable host: %+v %v", s, ok)
	}
	if _, ok := tbl.Get("zzz"); ok {
		t.Errorf("unknown host must report ok=false")
	}
}

func TestAllIsACopy(t *testing.T) {
	a := newFake(t, "m")
	tbl := health.New(map[string]string{"a": a.srv.URL}, time.Hour, nil)
	tbl.PollOnce(context.Background())
	all := tbl.All()
	all["a"].Loaded[0] = "changed"
	if s, _ := tbl.Get("a"); s.Loaded[0] != "m" {
		t.Errorf("All must return copies")
	}
	if len(all) != 1 {
		t.Errorf("All = %v", all)
	}
}

func TestRunPollsOnStart(t *testing.T) {
	a := newFake(t, "m")
	tbl := health.New(map[string]string{"a": a.srv.URL}, 20*time.Millisecond, nil)
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { tbl.Run(ctx); close(done) }()
	deadline := time.Now().Add(2 * time.Second)
	for a.hits.Load() < 3 && time.Now().Before(deadline) {
		time.Sleep(5 * time.Millisecond)
	}
	cancel()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("Run did not return after cancel")
	}
	if a.hits.Load() < 3 {
		t.Errorf("Run polled %d times in 2s at 20ms interval", a.hits.Load())
	}
	if s, _ := tbl.Get("a"); !s.Healthy {
		t.Errorf("not healthy after Run: %+v", s)
	}
}