webui4frpc/internal/cluster/balance_test.go

package cluster

import (
	"context"
	"encoding/json"
	"errors"
	"strconv"
	"strings"
	"testing"
	"time"

	"webui4frpc/internal/store"
)

// mustState round-trips the seed through JSON so the engine under test owns its
// own copy of Nodes/Topology. AdoptState assigns the state by value, but Nodes is
// a slice: handing one seed to several engines would alias one backing array and
// make a mutation on one node visible on the others — a test artefact that
// previously produced a false green here.
func mustState(t *testing.T, s State) *State {
	t.Helper()
	b, err := json.Marshal(s)
	if err != nil {
		t.Fatal(err)
	}
	var out State
	if err := json.Unmarshal(b, &out); err != nil {
		t.Fatal(err)
	}
	e := newTestEngine("probe", true)
	e.AdoptState(out)
	return e.State()
}

// migrateRing builds a 3-node ring where n1 owns every forward, plus n2/n3 idle.
func migrateRing(t *testing.T, ownCount int) (*Engine, *Engine, *Engine) {
	t.Helper()
	mk := func(id string, leader bool) *Engine {
		h := &fakeHandler{load: Load{MemPct: 5, NetPct: 5}}
		return NewEngine(id, id+":7500", "u", "p", "0.1.3", nil, h,
			func(ctx context.Context, next string, tk *Token) error { return nil },
			id+":7500", leader, "")
	}
	n1, n2, n3 := mk("n1", true), mk("n2", false), mk("n3", false)

	seed := State{
		LeaderID: "n1",
		Nodes: []Node{
			{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true, Load: Load{Forwards: ownCount, MemPct: 5}},
			{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{MemPct: 5}},
			{ID: "n3", Addr: "n3:7500", Alive: true, Load: Load{MemPct: 5}},
		},
		Topology:     map[string]*TopoEntry{},
		PendingTasks: map[string]*Task{},
	}
	for _, e := range []*Engine{n1, n2, n3} {
		var own State
		b, _ := json.Marshal(seed)
		if err := json.Unmarshal(b, &own); err != nil {
			t.Fatal(err)
		}
		e.AdoptState(own)
	}
	// n1 claims `ownCount` forwards.
	for i := 0; i < ownCount; i++ {
		n1.state.AddTopology(&Task{
			ID:    "f" + string(rune('a'+i)),
			Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
			Link: store.Link{RemotePort: 100 + i},
		}, "n1")
	}
	var snap State
	b, _ := json.Marshal(n1.state)
	_ = json.Unmarshal(b, &snap)
	for _, e := range []*Engine{n1, n2, n3} {
		e.AdoptState(snap)
	}
	return n1, n2, n3
}

// errNotWritable stands in for a failing settings write.
var errNotWritable = errors.New("not writable")

// ---- helpers ----

// balanceRing builds a 3-node ring with `owned` forwards on n1 (the evaluation
// node) and `others` on each of the other two, so a test can set up any
// distribution directly. The replicated Nodes[].Load is made consistent with the
// topology, because that is what a real token round leaves behind.
func balanceRing(t *testing.T, owned, others int) (*Engine, *Engine, *Engine) {
	t.Helper()
	mk := func(id string, leader bool) *Engine {
		h := &fakeHandler{load: Load{MemPct: 5, NetPct: 5}}
		return NewEngine(id, id+":7500", "u", "p", "0.1.3", nil, h,
			func(ctx context.Context, next string, tk *Token) error { return nil },
			id+":7500", leader, "")
	}
	n1, n2, n3 := mk("n1", true), mk("n2", false), mk("n3", false)

	seed := State{
		LeaderID: "n1",
		Nodes: []Node{
			{ID: "n1", Addr: "n1:7500", Alive: true, IsLeader: true, Load: Load{Forwards: owned, MemPct: 5}},
			{ID: "n2", Addr: "n2:7500", Alive: true, Load: Load{Forwards: others, MemPct: 5}},
			{ID: "n3", Addr: "n3:7500", Alive: true, Load: Load{Forwards: others, MemPct: 5}},
		},
		Topology:     map[string]*TopoEntry{},
		PendingTasks: map[string]*Task{},
	}
	for _, e := range []*Engine{n1, n2, n3} {
		var own State
		b, _ := json.Marshal(seed)
		if err := json.Unmarshal(b, &own); err != nil {
			t.Fatal(err)
		}
		e.AdoptState(own)
	}
	// Build the whole topology in n1's state, then broadcast it. Adding each
	// node's forwards to its OWN state and then re-adopting from n1 would discard
	// the others (one state is the source of the broadcast), which is how an
	// earlier version of this helper silently produced a 1-1-1 split regardless
	// of the arguments.
	add := func(owner, prefix string, count, basePort int) {
		for i := 0; i < count; i++ {
			n1.state.AddTopology(&Task{
				ID:    prefix + strconv.Itoa(i),
				Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
				Link: store.Link{RemotePort: basePort + i},
			}, owner)
		}
	}
	add("n1", "f", owned, 100)
	add("n2", "g", others, 200)
	add("n3", "h", others, 300)

	var snap State
	b, _ := json.Marshal(n1.state)
	_ = json.Unmarshal(b, &snap)
	for _, e := range []*Engine{n1, n2, n3} {
		e.AdoptState(snap)
	}
	return n1, n2, n3
}

// ---- the rule ----

// TestShareIsTheEvenSplit pins the arithmetic, including the ceiling. An even
// split of 3 over 3 is 1 (not 1.0 floored to 0): with a floor, a node holding
// one of three forwards would still be "over its share" and shed it.
func TestShareIsTheEvenSplit(t *testing.T) {
	for _, tc := range []struct{ total, alive, want int }{
		{3, 3, 1}, {4, 3, 2}, {6, 2, 3}, {1, 3, 1}, {0, 3, 0}, {7, 3, 3},
	} {
		got := (tc.total + tc.alive - 1) / tc.alive
		if got != tc.want {
			t.Errorf("share(%d over %d) = %d, want %d", tc.total, tc.alive, got, tc.want)
		}
	}
}

// TestOwnerReleasesWhenOverItsShare: 3 forwards on n1, none elsewhere → n1 holds
// 3 against a share of 1 and must give one up.
func TestOwnerReleasesWhenOverItsShare(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)

	pick, p, blocked := b.assess(1000)
	if pick == nil {
		t.Fatalf("n1 holds 3 of 3 forwards and should release one (blocked=%q)", blocked)
	}
	if blocked != "" {
		t.Errorf("nothing should block the first release, got %q", blocked)
	}
	if p.mine != 3 || p.share != 1 || p.total != 3 || p.alive != 3 {
		t.Errorf("pool = %+v, want mine=3 share=1 total=3 alive=3", p)
	}
	if pick.OwnerID != "n1" {
		t.Errorf("must release a forward n1 owns, got owner %q", pick.OwnerID)
	}
}

// TestReleasingStopsAtTheShare is the bound that keeps the rule from emptying a
// node: after shedding one, n1 holds 1 of 3 over 3 nodes, which is no longer more
// than the share, and it must stop.
func TestReleasingStopsAtTheShare(t *testing.T) {
	n1, _, _ := balanceRing(t, 1, 1) // even 1-1-1 split
	b := NewBalancer(n1, 0)

	pick, p, blocked := b.assess(1000)
	if pick != nil {
		t.Errorf("an even 1-1-1 split must not release anything, got %s:%d (blocked=%q)",
			pick.Local.Name, pick.Link.RemotePort, blocked)
	}
	if p.mine != 1 || p.share != 1 {
		t.Errorf("pool = %+v, want mine=1 share=1", p)
	}
}

// TestSingleNodeNeverReleases: a standalone node is the whole cluster, so its
// share is everything it holds. Without this the node would shed forwards onto
// nobody (they would sit unclaimed until it re-claimed them).
func TestSingleNodeNeverReleases(t *testing.T) {
	e := newTestEngine("solo", true)
	e.state.Nodes = []Node{{ID: "solo", Addr: "solo:7500", Alive: true, IsLeader: true}}
	for i := 0; i < 5; i++ {
		e.state.AddTopology(&Task{
			ID: "f" + strconv.Itoa(i), Local: store.Local{Name: "web"},
			Remote: store.Remote{Name: "frps1"}, Link: store.Link{RemotePort: 100 + i},
		}, "solo")
	}
	b := NewBalancer(e, 0)
	if pick, _, blocked := b.assess(1000); pick != nil {
		t.Errorf("a single-node ring must not release; got blocked=%q", blocked)
	}
}

// TestBelowFloorNeverReleases: two forwards on a node is the floor. A node
// holding exactly two at an even split of two must keep both — one release per
// node per pass is the granularity, and shedding below the floor would empty it.
func TestBelowFloorNeverReleases(t *testing.T) {
	n1, _, _ := balanceRing(t, 2, 2) // 6 total over 3 nodes → share 2
	b := NewBalancer(n1, 0)
	if pick, p, blocked := b.assess(1000); pick != nil {
		t.Errorf("2 owned against share %d must not release (blocked=%q, mine=%d)",
			p.share, blocked, p.mine)
	}
}

// TestNeverReleasesAStoppedForward: a stopped forward's entry is the carrier of
// the user's stop decision. Releasing it would hand another node a STOPPED
// forward as if it were work to run.
func TestNeverReleasesAStoppedForward(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	// Stop two of n1's three; the only live one must be chosen.
	for _, te := range n1.state.ForwardsOwnedBy("n1") {
		if te.Link.RemotePort == 100 || te.Link.RemotePort == 101 {
			n1.state.UpdateTopologyDisabled("web", "frps1", te.Link.RemotePort, true)
		}
	}
	b := NewBalancer(n1, 0)
	pick, _, _ := b.assess(1000)
	if pick == nil {
		t.Fatal("n1 still owns one live forward and is over its share")
	}
	if pick.Link.RemotePort != 102 {
		t.Errorf("released port %d; a stopped forward (100/101) must never be chosen",
			pick.Link.RemotePort)
	}

	// And with every forward stopped there is nothing to release.
	n1.state.UpdateTopologyDisabled("web", "frps1", 102, true)
	if pick, _, blocked := b.assess(1000); pick != nil {
		t.Errorf("all forwards stopped: nothing may be released (blocked=%q)", blocked)
	}
}

// TestCandidateChoiceIsDeterministic: the same state must pick the same forward,
// so the decision is explainable and a re-evaluation does not churn a different
// connection each tick.
func TestCandidateChoiceIsDeterministic(t *testing.T) {
	n1, _, _ := balanceRing(t, 4, 0)
	a := releaseCandidate(n1.State(), "n1")
	c := releaseCandidate(n1.State(), "n1")
	if a == nil || c == nil {
		t.Fatal("expected a candidate")
	}
	if a.TaskID != c.TaskID || a.Link.RemotePort != c.Link.RemotePort {
		t.Errorf("candidate changed between identical evaluations: %s:%d vs %s:%d",
			a.TaskID, a.Link.RemotePort, c.TaskID, c.Link.RemotePort)
	}
	if a.Link.RemotePort != 100 {
		t.Errorf("expected the lowest port (100), got %d", a.Link.RemotePort)
	}
}

// ---- the runtime signal that makes this safe ----

// TestRevokeLeavesACarrierAndReleaseRemovesIt is the reason the release path has
// to be a distinct task kind.
//
// A user stop keeps the topology entry as the carrier of the disabled flag, so
// the entry is STILL THERE and still attributed to the owner. A rebalance is the
// opposite: the entry must go, because (a) the node's load is derived from
// owned entries so it cannot fall while the entry exists, and (b) every node's
// duplicate guard refuses a creation task for a forward that already has an owner
// — so the re-published task would be dropped by everyone and the forward
// stranded on the node that gave it up.
func TestRevokeLeavesACarrierAndReleaseRemovesIt(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)

	// A stop marks the entry and keeps it.
	n1.state.UpdateTopologyDisabled("web", "frps1", 100, true)
	tk := &Task{Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
		Link: store.Link{RemotePort: 100}}
	if owner := n1.state.TopologyOwner(tk); owner != "n1" {
		t.Fatalf("a stop must keep the entry as the flag's carrier, owner=%q", owner)
	}
	if got := len(n1.state.ForwardsOwnedBy("n1")); got != 3 {
		t.Fatalf("a stop must not change the owner's forward count, got %d", got)
	}

	// A release removes the entry (leaving the other two alone).
	if !n1.RemoveTopologyEntry("web", "frps1", 100) {
		t.Fatal("release must remove the entry")
	}
	if owner := n1.state.TopologyOwner(tk); owner != "" {
		t.Errorf("after a release the forward must have no owner, got %q", owner)
	}
	if got := len(n1.state.ForwardsOwnedBy("n1")); got != 2 {
		t.Errorf("a release must reduce the owner's count, got %d", got)
	}
	if got := len(n1.state.TopologyList()); got != 2 {
		t.Errorf("the other forwards must be untouched, got %d entries", got)
	}
}

// TestReleaseTaskExecutesOnlyOnTheOwner: like a revoke, a release is
// owner-directed. Executing it elsewhere would be an attempt to move somebody
// else's load, which cannot work — the owner's own count is what must fall.
func TestReleaseTaskExecutesOnlyOnTheOwner(t *testing.T) {
	n1, n2, _ := balanceRing(t, 3, 0)

	released := map[string]int{}
	for _, e := range []*Engine{n1, n2} {
		h := e.Handler.(*fakeHandler)
		h.release = func(ctx context.Context, tk *Task) error {
			released[e.ID]++
			return nil
		}
	}
	tk := n1.ReleaseTask(store.Local{Name: "web"}, store.Remote{Name: "frps1"},
		store.Link{RemotePort: 100})
	if tk == nil {
		t.Fatal("ReleaseTask returned nil")
	}

	// Each Engine keeps its OWN state, so a task published on n1 is not visible
	// to n2 until a token round carries it. Broadcast it the way the ring does.
	var carry State
	cb, _ := json.Marshal(n1.state)
	_ = json.Unmarshal(cb, &carry)
	n2.AdoptState(carry)
	n2.state.PendingTasks[tk.ID] = tk

	// n2 is the lightest node and does NOT own the forward: it must leave the
	// task alone.
	if err := n2.runCommands(context.Background(), &Token{}); err != nil {
		t.Fatalf("n2 runCommands: %v", err)
	}
	if released["n2"] != 0 {
		t.Errorf("n2 released %d forward(s); a release is owner-directed only", released["n2"])
	}
	if _, still := n2.state.PendingTasks[tk.ID]; !still {
		t.Error("n2 consumed a release it does not own; it must ride the token to the owner")
	}

	// n1 owns it: it must release, drop the entry, and re-publish the forward.
	if err := n1.runCommands(context.Background(), &Token{}); err != nil {
		t.Fatalf("n1 runCommands: %v", err)
	}
	if released["n1"] != 1 {
		t.Fatalf("the owner released %d times, want 1", released["n1"])
	}
	if owner := n1.state.TopologyOwner(&Task{
		Local: store.Local{Name: "web"}, Remote: store.Remote{Name: "frps1"},
		Link: store.Link{RemotePort: 100},
	}); owner != "" {
		t.Errorf("the released forward still has owner %q", owner)
	}
	// And it is back in the ring as an ordinary creation task.
	var republished *Task
	for _, p := range n1.state.PendingList() {
		if !p.Release && !p.Revoke && !p.Restart &&
			p.Link.RemotePort == 100 && p.Local.Name == "web" {
			republished = p
		}
	}
	if republished == nil {
		t.Fatal("the released forward was not re-published as an unclaimed task")
	}
	if republished.Link.Disabled {
		t.Error("the re-published forward must be ENABLED, not marked stopped")
	}
}

// ---- dry-run / off ----

// TestDryRunReleasesNothing: dry-run is the default, and its whole value is that
// it can run on a live cluster without disturbing anything.
func TestDryRunReleasesNothing(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceDryRun)

	msg := b.Step(1000)
	if !strings.Contains(msg, "would release") {
		t.Fatalf("dry-run should report what it would do, got %q", msg)
	}
	if b.InFlight() {
		t.Error("dry-run started a release")
	}
	for _, tk := range n1.state.PendingList() {
		if tk.Release {
			t.Error("dry-run published a release task")
		}
	}
	if got := len(n1.state.TopologyList()); got != 3 {
		t.Errorf("dry-run changed the topology: %d entries, want 3", got)
	}
	// A second identical tick must not repeat the same line (log hygiene).
	if again := b.Step(1010); again != "" {
		t.Errorf("dry-run repeated an unchanged decision: %q", again)
	}
}

// TestModeOffDoesNothing reinforces that OFF means off, not "dry-run".
func TestModeOffDoesNothing(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceOff)
	if msg := b.Step(1000); msg != "" {
		t.Errorf("mode off reported %q", msg)
	}
	if b.InFlight() {
		t.Error("mode off started a release")
	}
	// A typo must not enable releasing.
	b.SetMode("yes-please")
	if b.Mode() != BalanceDryRun {
		t.Errorf("an unknown mode must fall back to dry-run, got %q", b.Mode())
	}
}

// ---- end-to-end through the ring machinery ----

// TestOwnerReleaseThenRingClaimsElsewhere is the whole design in one test: the
// owner gives the forward up, and the ordinary lowest-load claim then places it
// on an idle node — exactly what the self-directed rule depends on.
func TestOwnerReleaseThenRingClaimsElsewhere(t *testing.T) {
	n1, n2, n3 := balanceRing(t, 3, 0)

	// n2 and n3 must look idle to the claim path (they do: Forwards=0).
	claimed := map[string]int{}
	for _, e := range []*Engine{n1, n2, n3} {
		h := e.Handler.(*fakeHandler)
		h.claim = func(ctx context.Context, tk *Task) error {
			claimed[e.ID]++
			return nil
		}
	}
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceOn)

	// n1 is over its share → it releases.
	msg := b.Step(1000)
	if !strings.Contains(msg, "releasing") {
		t.Fatalf("expected a release, got %q", msg)
	}
	if err := n1.runCommands(context.Background(), &Token{}); err != nil {
		t.Fatalf("n1 runCommands: %v", err)
	}
	// Find the task the owner re-published for the forward it gave up.
	var tk *Task
	for _, p := range n1.state.PendingList() {
		if !p.Release && p.Link.RemotePort == 100 && p.Local.Name == "web" {
			tk = p
		}
	}
	if tk == nil {
		t.Fatal("the owner did not re-publish the forward it released")
	}
	if got := len(n1.state.ForwardsOwnedBy("n1")); got != 2 {
		t.Fatalf("n1 should now own 2, got %d", got)
	}

	// The forward is unclaimed; the lowest node claims it. n2 is the lightest
	// (Forwards=0 and the same mem/net), so it should take it. The task reaches
	// n2 on the token — each engine holds its own state, so carry it over the way
	// a real round does.
	var carried State
	cb, _ := json.Marshal(n1.state)
	_ = json.Unmarshal(cb, &carried)
	n2.AdoptState(carried)
	n2.state.PendingTasks[tk.ID] = tk

	if err := n2.runCommands(context.Background(), &Token{}); err != nil {
		t.Fatalf("n2 runCommands: %v", err)
	}
	if claimed["n2"] == 0 {
		t.Fatal("the idle node did not claim the released forward")
	}
	// And n1 no longer owns it.
	for _, te := range n1.state.ForwardsOwnedBy("n1") {
		if te.Link.RemotePort == 100 {
			t.Error("n1 still owns the forward it released")
		}
	}

	// The balancer sees the release land and stops holding the release state.
	if msg := b.Step(1100); !strings.Contains(msg, "gave up") {
		t.Errorf("expected the release to be confirmed, got %q", msg)
	}
	if b.InFlight() {
		t.Error("the release should be complete")
	}
}

// TestCooldownBoundsTheChurn: a release costs a live connection, so one node may
// not shed again until the cooldown elapses — even if it is still over its share.
func TestCooldownBoundsTheChurn(t *testing.T) {
	n1, _, _ := balanceRing(t, 4, 0)
	b := NewBalancer(n1, 600) // 10 minutes
	b.SetMode(BalanceOn)

	if msg := b.Step(1000); !strings.Contains(msg, "releasing") {
		t.Fatalf("first release expected, got %q", msg)
	}
	n1.RemoveTopologyEntry("web", "frps1", 100) // land it
	if msg := b.Step(1010); !strings.Contains(msg, "gave up") {
		t.Fatalf("expected confirmation, got %q", msg)
	}
	// Still over its share (3 of 4 over 3 → share 2), but inside the cooldown.
	if pick, _, blocked := b.assess(1020); pick == nil || blocked != "cooldown" {
		t.Errorf("expected the second release to be blocked by cooldown, got blocked=%q", blocked)
	}
	if b.InFlight() {
		t.Error("cooldown must not leave a release in flight")
	}
	// Past the cooldown it may act again.
	if pick, _, blocked := b.assess(1000 + 601); pick == nil || blocked != "" {
		t.Errorf("past the cooldown the release should be allowed, got blocked=%q", blocked)
	}
}

// TestReleaseDeadlineClearsAStuckConfirmation: if the release never lands (leader
// change, ring join), the node must not wedge forever.
func TestReleaseDeadlineClearsAStuckConfirmation(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceOn)
	_ = b.Step(1000) // release published, entry NOT removed (nothing executed it)
	if !b.InFlight() {
		t.Fatal("precondition: a release should be in flight")
	}
	if msg := b.Step(1000 + releaseDeadlineSec + 1); msg != "" {
		t.Errorf("a timed-out release should clear silently, got %q", msg)
	}
	if b.InFlight() {
		t.Error("the release state should have been cleared after the deadline")
	}
}

// ---- read-only report ----

// TestReportDoesNotMutate — polling the endpoint must not start, delay or cancel
// a release. A read-only surface that perturbs what it reports is worse than none.
func TestReportDoesNotMutate(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 600)
	b.SetMode(BalanceOn)

	for i := 0; i < 5; i++ {
		rep := b.Report()
		if rep.SelfForwards != 3 {
			t.Fatalf("report changed the node's forward count: %d", rep.SelfForwards)
		}
		if rep.Mode != BalanceOn {
			t.Fatalf("report changed the mode: %q", rep.Mode)
		}
	}
	if b.InFlight() || b.lastReleaseAt != 0 {
		t.Error("polling the report started a release")
	}
	if rep := b.Report(); rep.WouldRelease == nil {
		t.Error("the report should still surface the pending decision")
	}
}

// TestReportExplainsWhyNothingHappens pins the field that answers the question
// the log cannot: an over-share node that is holding back must say why.
func TestReportExplainsWhyNothingHappens(t *testing.T) {
	n1, _, _ := balanceRing(t, 4, 0)
	b := NewBalancer(n1, 600)
	b.SetMode(BalanceOn)

	rep := b.Report()
	if !rep.SelfOverShare {
		t.Fatalf("4 of 4 forwards over 3 nodes should be over the share (share=%d)", rep.Share)
	}
	if rep.WouldRelease == nil {
		t.Fatal("an over-share node must surface the forward it would give up")
	}
	if rep.BlockedBy != "" {
		t.Errorf("nothing should block the first release, got %q", rep.BlockedBy)
	}

	// Inside the cooldown the same node reports the block. Report() reads the
	// real clock, so set lastReleaseAt relative to now rather than to a fake one.
	b.lastReleaseAt = time.Now().Unix() - 10
	n1.RemoveTopologyEntry("web", "frps1", 100)
	rep = b.Report()
	if rep.BlockedBy != "cooldown" {
		t.Errorf("after a release the node should report a cooldown block, got %q", rep.BlockedBy)
	}
}

// TestReportOnBalancedNodeIsQuiet: the healthy case must be a real answer, not an
// empty one — otherwise "balanced" and "balancing is switched off" look identical.
func TestReportOnBalancedNodeIsQuiet(t *testing.T) {
	n1, _, _ := balanceRing(t, 1, 1)
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceOn)

	rep := b.Report()
	if rep.WouldRelease != nil {
		t.Errorf("a balanced node reported a release: %+v", rep.WouldRelease)
	}
	if rep.SelfOverShare {
		t.Error("1 of 3 over 3 nodes is not over the share")
	}
	if rep.TotalForwards != 3 || rep.AliveNodes != 3 || rep.Share != 1 {
		t.Errorf("report pool = total %d alive %d share %d, want 3/3/1",
			rep.TotalForwards, rep.AliveNodes, rep.Share)
	}
}

// ---- mode switching (the cluster page's switch) ----

// TestStoredModeWinsOverEnv: flipping the switch must not silently revert on the
// next restart. The environment is only the bootstrap default; once an operator
// has chosen a mode it is persisted and wins.
func TestStoredModeWinsOverEnv(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode("dry-run") // the env default

	var saved string
	b.SetModePersistence(
		func() (string, bool) { return "on", true },
		func(m string) error { saved = m; return nil },
	)
	if b.Mode() != BalanceOn {
		t.Fatalf("a stored mode must win over the environment default, got %q", b.Mode())
	}
	if b.ModeSource() != "stored" {
		t.Errorf("ModeSource = %q, want stored", b.ModeSource())
	}
	_ = saved
}

// TestUnconfiguredFallsBackToEnv: with nothing stored, the environment default
// stands — otherwise every node would need a settings row to behave as documented.
func TestUnconfiguredFallsBackToEnv(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode("off")
	b.SetModePersistence(func() (string, bool) { return "", false }, func(string) error { return nil })
	if b.Mode() != BalanceOff {
		t.Errorf("with nothing stored the env default must stand, got %q", b.Mode())
	}
	if b.ModeSource() != "stored" {
		// persist is installed, so the node CAN be controlled
		t.Logf("ModeSource=%q", b.ModeSource())
	}
}

// TestPersistBeforeApply: if the write fails, the mode must not change.
// Applying first would leave the node running a mode it cannot reproduce after a
// restart — the switch would lie.
func TestPersistBeforeApply(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceDryRun)
	b.SetModePersistence(nil, func(string) error { return errNotWritable })

	if err := b.SetModeFromUI(BalanceOn); err == nil {
		t.Fatal("a failed save must surface an error")
	}
	if b.Mode() != BalanceDryRun {
		t.Errorf("a failed save must leave the mode unchanged, got %q", b.Mode())
	}
}

// TestSetModeFromUIRejectsNonsense: the endpoint takes a string from the wire, so
// it must refuse anything it does not understand rather than guessing.
func TestSetModeFromUIRejectsNonsense(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceDryRun)
	if err := b.SetModeFromUI("yes"); err == nil {
		t.Fatal("an unknown mode must be rejected, not coerced")
	}
	if b.Mode() != BalanceDryRun {
		t.Errorf("a rejected mode must not change anything, got %q", b.Mode())
	}
}

// TestSetModeFromUIPersistsAndApplies: the happy path, both effects.
func TestSetModeFromUIPersistsAndApplies(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	b.SetMode(BalanceDryRun)
	var saved string
	b.SetModePersistence(nil, func(m string) error { saved = m; return nil })

	if err := b.SetModeFromUI(BalanceOn); err != nil {
		t.Fatalf("SetModeFromUI: %v", err)
	}
	if saved != BalanceOn {
		t.Errorf("saved %q, want on", saved)
	}
	if b.Mode() != BalanceOn {
		t.Errorf("mode = %q, want on", b.Mode())
	}
	if !b.Report().CanControl {
		t.Error("a node with persistence installed should report itself controllable")
	}
}

// TestReportCanControlOffWithoutPersistence: a node that cannot persist a choice
// must not advertise the switch (an embedder/test without a store).
func TestReportCanControlOffWithoutPersistence(t *testing.T) {
	n1, _, _ := balanceRing(t, 3, 0)
	b := NewBalancer(n1, 0)
	if b.Report().CanControl {
		t.Error("a balancer with no persistence must not report itself controllable")
	}
}