llmgw/internal/cli/ctl_upgrade_test.go

package cli

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

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

// fakeFleet is an admin API whose nodes come back on the version they are told
// to upgrade to — except those listed in stuck.
type fakeFleet struct {
	mu       sync.Mutex
	versions map[string]string
	stuck    map[string]bool
	told     []string
}

func (f *fakeFleet) server(t *testing.T) *httptest.Server {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		f.mu.Lock()
		defer f.mu.Unlock()
		switch {
		case r.URL.Path == "/api/v1/nodes":
			var out []adminapi.NodeView
			for id, v := range f.versions {
				out = append(out, adminapi.NodeView{ID: id, State: "ready", AgentVersion: v})
			}
			_ = json.NewEncoder(w).Encode(out)
		case r.URL.Path == "/api/v1/replicas":
			var out []adminapi.ReplicaView
			for id := range f.versions {
				out = append(out, adminapi.ReplicaView{ID: "r-" + id, State: "ready",
					Shards: []adminapi.ShardView{{Node: id}}})
			}
			_ = json.NewEncoder(w).Encode(out)
		case strings.HasSuffix(r.URL.Path, "/upgrade"):
			id := strings.TrimSuffix(strings.TrimPrefix(r.URL.Path, "/api/v1/nodes/"), "/upgrade")
			var req adminapi.UpgradeRequest
			_ = json.NewDecoder(r.Body).Decode(&req)
			f.told = append(f.told, id)
			if !f.stuck[id] {
				f.versions[id] = req.Version
			}
			_, _ = w.Write([]byte(`{}`))
		default:
			http.NotFound(w, r)
		}
	}))
	t.Cleanup(srv.Close)
	return srv
}

func TestAnUpgradeSkipsTheNodesAlreadyDone(t *testing.T) {
	f := &fakeFleet{versions: map[string]string{"a": "v1.0.0", "b": "v1.0.1", "c": "v1.0.0"}}
	srv := f.server(t)
	var out bytes.Buffer
	err := upgradeNodes(t.Context(), &out, func() *ctlClient { return newCtlClient(srv.URL) }, "v1.0.1", nil, 5*time.Second, 1)
	if err != nil {
		t.Fatal(err)
	}
	if strings.Join(f.told, ",") != "a,c" {
		t.Errorf("told %v; b was already on the version", f.told)
	}
	if !strings.Contains(out.String(), "2 node(s) upgraded to v1.0.1 in") || !strings.Contains(out.String(), "1 already on it") {
		t.Errorf("output:\n%s", out.String())
	}
}

func TestAnUpgradeStopsAtTheFirstWaveThatDoesNotComeBack(t *testing.T) {
	// A bad release reaches one node, not all of them.
	f := &fakeFleet{versions: map[string]string{"a": "v1.0.0", "b": "v1.0.0"}, stuck: map[string]bool{"a": true}}
	srv := f.server(t)
	var out bytes.Buffer
	err := upgradeNodes(t.Context(), &out, func() *ctlClient { return newCtlClient(srv.URL) }, "v1.0.1", nil, 3*time.Second, 1)
	if err == nil || !strings.Contains(err.Error(), "not back within") || !strings.Contains(err.Error(), "a (version") {
		t.Fatalf("err %v", err)
	}
	if len(f.told) != 1 {
		t.Errorf("went on to %v after a node failed to come back", f.told)
	}
}

func TestAnUpgradeGoesInWavesThatGrowUpToTheLimit(t *testing.T) {
	// One first, then twice as many each time, up to --parallel: hundreds of
	// nodes in minutes rather than hours, and a bad release still on one.
	versions := map[string]string{}
	for i := range 12 {
		versions[fmt.Sprintf("n%02d", i)] = "v1.0.0"
	}
	f := &fakeFleet{versions: versions}
	srv := f.server(t)
	var out bytes.Buffer
	if err := upgradeNodes(t.Context(), &out, func() *ctlClient { return newCtlClient(srv.URL) }, "v1.0.1", nil, 5*time.Second, 4); err != nil {
		t.Fatal(err)
	}
	var sizes []string
	for _, line := range strings.Split(out.String(), "\n") {
		if n, ok := strings.CutPrefix(line, "== "); ok {
			sizes = append(sizes, strings.Fields(n)[0])
		}
	}
	if got := strings.Join(sizes, ","); got != "1,2,4,4,1" {
		t.Errorf("waves %s:\n%s", got, out.String())
	}
}

func TestAnUpgradeKnowsWhetherItRunsOnTheLeader(t *testing.T) {
	sum := adminapi.Summary{}
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
		_ = json.NewEncoder(w).Encode(sum)
	}))
	defer srv.Close()
	client := func() *ctlClient { return newCtlClient(srv.URL) }
	if !thisLeads(t.Context(), client) {
		t.Error("a control plane with no rival does not lead")
	}
	sum.ControlPlanes.Members = []adminapi.ControlPlane{{Name: "a", State: "leading"}, {Name: "b", State: "standing by", This: true}}
	if thisLeads(t.Context(), client) {
		t.Error("a standby took itself for the leader: it would upgrade the leader before itself")
	}
	sum.ControlPlanes.Members[0].This, sum.ControlPlanes.Members[1].This = true, false
	if !thisLeads(t.Context(), client) {
		t.Error("the leader took itself for a standby")
	}
}