pangolin/server/internal/scheduler/scheduler_test.go

package scheduler_test

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/alicebob/miniredis/v2"
	"github.com/redis/go-redis/v9"

	"github.com/wangjia/pangolin/server/internal/scheduler"
	"github.com/wangjia/pangolin/server/internal/scheduler/detect"
	"github.com/wangjia/pangolin/server/internal/scheduler/orchestrate"
	"github.com/wangjia/pangolin/server/internal/scheduler/probe"
)

// ─────────────────────────────────────────────────────────────────────────────
// Shared test helpers
// ─────────────────────────────────────────────────────────────────────────────

func newTestRedis(t *testing.T) (*redis.Client, *miniredis.Miniredis) {
	t.Helper()
	mr := miniredis.RunT(t)
	rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
	return rdb, mr
}

// testIntervals returns Config timing values suitable for fast unit tests.
func testIntervals() (leaseTTL, renew, retry, tick, interval time.Duration) {
	return 200 * time.Millisecond, // LeaseTTL
		50 * time.Millisecond,  // RenewPeriod
		50 * time.Millisecond,  // RetryPeriod
		5 * time.Second,        // TickTimeout
		80 * time.Millisecond   // loop interval
}

// counterTick returns a tick function that atomically increments *n.
func counterTick(n *int64) func(context.Context) error {
	return func(_ context.Context) error {
		atomic.AddInt64(n, 1)
		return nil
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Stub implementations for scheduler.Config fields
// ─────────────────────────────────────────────────────────────────────────────

// stubEngine implements scheduler.DetectEngine.
type stubEngine struct{ count int64 }

func (e *stubEngine) Tick(_ context.Context) error { atomic.AddInt64(&e.count, 1); return nil }

// stubReplacer implements scheduler.OrchestrateReplacer.
type stubReplacer struct{ count int64 }

func (r *stubReplacer) Tick(_ context.Context) error { atomic.AddInt64(&r.count, 1); return nil }

// stubGrayscale implements scheduler.OrchestrateGrayscale.
type stubGrayscale struct{ count int64 }

func (g *stubGrayscale) Advance(_ context.Context) error {
	atomic.AddInt64(&g.count, 1)
	return nil
}

// stubProbeStore implements scheduler.ProbeStateReader.
type stubProbeStore struct {
	mu    sync.Mutex
	alive []string
}

func (s *stubProbeStore) AliveProbes(_ context.Context) ([]string, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return append([]string{}, s.alive...), nil
}

func (s *stubProbeStore) setAlive(ids ...string) {
	s.mu.Lock(); defer s.mu.Unlock()
	s.alive = ids
}

// recordingNotifier captures NotifyFault calls.
type recordingNotifier struct {
	mu     sync.Mutex
	calls  []string // nodeID values
}

func (r *recordingNotifier) NotifyFault(_ context.Context, nodeID, _ string) error {
	r.mu.Lock(); defer r.mu.Unlock()
	r.calls = append(r.calls, nodeID)
	return nil
}

func (r *recordingNotifier) callCount() int {
	r.mu.Lock(); defer r.mu.Unlock()
	return len(r.calls)
}

func (r *recordingNotifier) called(nodeID string) bool {
	r.mu.Lock(); defer r.mu.Unlock()
	for _, id := range r.calls {
		if id == nodeID { return true }
	}
	return false
}

// ─────────────────────────────────────────────────────────────────────────────
// TestLeaderElection: only one of two competing instances runs ticks
// ─────────────────────────────────────────────────────────────────────────────

func TestLeaderElection(t *testing.T) {
	rdb, _ := newTestRedis(t)
	leaseTTL, renew, retry, tickTout, interval := testIntervals()

	var tickA, tickB int64
	makeEngine := func(n *int64) *stubEngine {
		return &stubEngine{}
	}
	_ = makeEngine

	buildCfg := func(id string, engine *stubEngine, replacer *stubReplacer, gray *stubGrayscale) scheduler.Config {
		return scheduler.Config{
			RDB:                 rdb,
			InstanceID:          id,
			Engine:              engine,
			Replacer:            replacer,
			Grayscale:           gray,
			DetectInterval:      interval,
			OrchestrateInterval: interval,
			CapacityInterval:    interval,
			LeaseTTL:            leaseTTL,
			RenewPeriod:         renew,
			RetryPeriod:         retry,
			TickTimeout:         tickTout,
		}
	}

	engA := &stubEngine{}
	repA := &stubReplacer{}
	gryA := &stubGrayscale{}
	schedA := scheduler.New(buildCfg("inst-A", engA, repA, gryA))

	engB := &stubEngine{}
	repB := &stubReplacer{}
	gryB := &stubGrayscale{}
	schedB := scheduler.New(buildCfg("inst-B", engB, repB, gryB))

	ctxA, cancelA := context.WithCancel(context.Background())
	ctxB, cancelB := context.WithCancel(context.Background())
	defer cancelA()
	defer cancelB()

	// Run both schedulers; they compete for "detect" / "orchestrate" / "capacity" leases.
	var wgA, wgB sync.WaitGroup
	wgA.Add(1)
	go func() { defer wgA.Done(); _ = schedA.Start(ctxA) }()
	wgB.Add(1)
	go func() { defer wgB.Done(); _ = schedB.Start(ctxB) }()

	// Let them run for a few tick intervals.
	time.Sleep(500 * time.Millisecond)

	// Collect tick counts.
	tickA = atomic.LoadInt64(&engA.count)
	tickB = atomic.LoadInt64(&engB.count)
	t.Logf("after 500ms: tickA=%d tickB=%d", tickA, tickB)

	// At most one instance should have run detect ticks.
	if tickA > 0 && tickB > 0 {
		t.Errorf("both inst-A and inst-B ran detect ticks — leader election broken (A=%d B=%d)", tickA, tickB)
	}
	if tickA == 0 && tickB == 0 {
		t.Error("neither instance ran any detect ticks — scheduler not ticking")
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// TestFollowerTakeover: follower takes over after leader's context is cancelled
// ─────────────────────────────────────────────────────────────────────────────

func TestFollowerTakeover(t *testing.T) {
	rdb, mr := newTestRedis(t)
	leaseTTL, renew, retry, tickTout, interval := testIntervals()

	buildCfg := func(id string) scheduler.Config {
		return scheduler.Config{
			RDB:                 rdb,
			InstanceID:          id,
			Engine:              &stubEngine{},
			Replacer:            &stubReplacer{},
			Grayscale:           &stubGrayscale{},
			DetectInterval:      interval,
			OrchestrateInterval: interval,
			CapacityInterval:    interval,
			LeaseTTL:            leaseTTL,
			RenewPeriod:         renew,
			RetryPeriod:         retry,
			TickTimeout:         tickTout,
		}
	}

	// Use a simple counting tick for "detect" loop wired via Engine.
	var ticksLeader, ticksFollower int64

	engLeader := &counterEngine{n: &ticksLeader}
	engFollower := &counterEngine{n: &ticksFollower}

	cfgLeader := buildCfg("leader")
	cfgLeader.Engine = engLeader
	cfgFollower := buildCfg("follower")
	cfgFollower.Engine = engFollower

	schedLeader := scheduler.New(cfgLeader)
	schedFollower := scheduler.New(cfgFollower)

	ctxLeader, cancelLeader := context.WithCancel(context.Background())
	ctxFollower, cancelFollower := context.WithCancel(context.Background())
	defer cancelFollower()

	var wgLeader sync.WaitGroup
	wgLeader.Add(1)
	go func() { defer wgLeader.Done(); _ = schedLeader.Start(ctxLeader) }()

	// Start ONLY the leader first and wait until it actually ticks, so it is the
	// confirmed owner of the detect lease before the follower joins. Starting both
	// instances at once races the leader election — either could win the lease —
	// which previously flaked this test as "leader never ticked".
	for i := 0; i < 150; i++ {
		if atomic.LoadInt64(&ticksLeader) > 0 {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	if atomic.LoadInt64(&ticksLeader) == 0 {
		t.Fatal("leader never ticked")
	}

	// Now start the follower; it blocks retrying to acquire the held leases.
	go func() { _ = schedFollower.Start(ctxFollower) }()

	prevFollower := atomic.LoadInt64(&ticksFollower)

	// Stop leader — its defer releases the leader key immediately.
	cancelLeader()
	wgLeader.Wait()

	// Fast-forward miniredis clock to expire the lease (belt-and-suspenders for
	// cases where release didn't fire, e.g. kill -9 simulation).
	mr.FastForward(leaseTTL + 10*time.Millisecond)

	// Follower should take over shortly after the lease frees (≈ RetryPeriod +
	// one loop interval ≈ 130ms). 3s deadline leaves generous headroom for
	// goroutine-scheduling jitter under full-suite CPU load; the assertion still
	// fails if takeover genuinely never happens.
	deadline := time.Now().Add(3 * time.Second)
	for time.Now().Before(deadline) {
		if atomic.LoadInt64(&ticksFollower) > prevFollower {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}

	if atomic.LoadInt64(&ticksFollower) <= prevFollower {
		t.Errorf("follower did not take over after leader stopped (leader=%d follower=%d→%d)",
			atomic.LoadInt64(&ticksLeader),
			prevFollower, atomic.LoadInt64(&ticksFollower))
	}
}

// counterEngine is an Engine whose Tick increments *n — used in takeover test.
type counterEngine struct{ n *int64 }

func (e *counterEngine) Tick(_ context.Context) error { atomic.AddInt64(e.n, 1); return nil }

// ─────────────────────────────────────────────────────────────────────────────
// TestGracefulShutdown: leader key is released on SIGTERM-equivalent ctx cancel
// ─────────────────────────────────────────────────────────────────────────────

func TestGracefulShutdown(t *testing.T) {
	rdb, _ := newTestRedis(t)
	leaseTTL, renew, retry, tickTout, interval := testIntervals()

	sched := scheduler.New(scheduler.Config{
		RDB:                 rdb,
		InstanceID:          "shutdown-test",
		Engine:              &stubEngine{},
		Replacer:            &stubReplacer{},
		Grayscale:           &stubGrayscale{},
		DetectInterval:      interval,
		OrchestrateInterval: interval,
		CapacityInterval:    interval,
		LeaseTTL:            leaseTTL,
		RenewPeriod:         renew,
		RetryPeriod:         retry,
		TickTimeout:         tickTout,
	})

	ctx, cancel := context.WithCancel(context.Background())

	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); _ = sched.Start(ctx) }()

	// Wait until at least one leader key is acquired.
	deadline := time.Now().Add(500 * time.Millisecond)
	for time.Now().Before(deadline) {
		v, _ := rdb.Exists(context.Background(), "sched:leader:detect").Result()
		if v > 0 {
			break
		}
		time.Sleep(10 * time.Millisecond)
	}

	// Signal shutdown — equivalent to SIGTERM.
	cancel()
	done := make(chan struct{})
	go func() { wg.Wait(); close(done) }()
	select {
	case <-done:
	case <-time.After(2 * time.Second):
		t.Fatal("scheduler did not stop within 2s")
	}

	// Verify all three leader keys are gone (released by the scheduler, not expired).
	ctx2 := context.Background()
	for _, loop := range []string{"detect", "orchestrate", "capacity"} {
		key := "sched:leader:" + loop
		v, err := rdb.Exists(ctx2, key).Result()
		if err != nil {
			t.Fatalf("EXISTS %s: %v", key, err)
		}
		if v != 0 {
			t.Errorf("leader key %q not released on shutdown", key)
		}
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// TestCapacityTickProbeDisconnect: disconnected probe triggers Notifier
// ─────────────────────────────────────────────────────────────────────────────

func TestCapacityTickProbeDisconnect(t *testing.T) {
	rdb, _ := newTestRedis(t)
	leaseTTL, renew, retry, tickTout, interval := testIntervals()

	probeStore := &stubProbeStore{}
	notifier := &recordingNotifier{}

	// "probe-sg-01" is known but absent from alive set → should fire alert.
	probeStore.setAlive("probe-jp-01") // only jp is alive

	sched := scheduler.New(scheduler.Config{
		RDB:                 rdb,
		InstanceID:          "capacity-test",
		Engine:              &stubEngine{},
		Replacer:            &stubReplacer{},
		Grayscale:           &stubGrayscale{},
		DetectInterval:      interval,
		OrchestrateInterval: interval,
		CapacityInterval:    interval,
		LeaseTTL:            leaseTTL,
		RenewPeriod:         renew,
		RetryPeriod:         retry,
		TickTimeout:         tickTout,
		ProbeStore:          probeStore,
		KnownProbeIDs:       []string{"probe-sg-01", "probe-jp-01"},
		Notifier:            notifier,
	})

	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); _ = sched.Start(ctx) }()

	// Wait for the capacity tick to fire at least once.
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		if notifier.callCount() > 0 {
			break
		}
		time.Sleep(20 * time.Millisecond)
	}
	cancel()
	wg.Wait()

	if !notifier.called("probe-sg-01") {
		t.Errorf("expected NotifyFault for probe-sg-01 (disconnected), but it was not called")
	}
	if notifier.called("probe-jp-01") {
		t.Errorf("unexpected NotifyFault for probe-jp-01 (it is alive)")
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// End-to-end mock scenario
//
// Verifies the complete pipeline: blocked node detection → replacement
// orchestration → new node activated → grayscale started.
//
// Mock wiring:
//   - #5  LifecycleService  → detect.MockLifecycle + mockOrchestrateLC
//   - #14 ProvisionService  → mockProvision
//   - probe store           → mockSnapshotter (inject fake probe data)
// ─────────────────────────────────────────────────────────────────────────────

const e2eNode = "node-sg-001"

func TestE2EMockScenario(t *testing.T) {
	ctx := context.Background()
	rdb, _ := newTestRedis(t)

	// ── 15D: detection engine ──────────────────────────────────────────────────

	detectLC := detect.NewMockLifecycle([]detect.NodeInfo{
		{ID: e2eNode, Status: detect.StatusUp, Weight: 100},
	})
	snapper := &mockSnapshotter{data: make(map[string]map[string]probe.ProbeSnapshot)}
	detectStreaks := detect.NewStreakStore(rdb)
	detectCfg := detect.DefaultConfig() // SuspectStreakMin=2, ConfirmedStreakMin=6
	engine := detect.NewEngine(snapper, detectLC, detectStreaks, rdb, nil, &detectCfg)

	// ── 15E: orchestration engine ──────────────────────────────────────────────

	orchLC := newMockOrchestrateLC()
	orchLC.addNode(&orchestrate.NodeInfo{
		ID:     e2eNode,
		Tier:   "premium",
		Region: "ap-southeast-1",
		Role:   "vpn",
	})
	prov := newMockProvision(orchestrate.ProviderInfo{ID: "vultr"})
	replacer := orchestrate.NewReplacer(orchestrate.Config{
		RDB:      rdb,
		Prov:     prov,
		LC:       orchLC,
		Snaps:    snapper, // probe.Store implements both ProbeSnapshotter interfaces
		Breaker:  orchestrate.StubBreaker{},
		Notifier: orchestrate.LogNotifier{},
		Clock:    orchestrate.RealClock{},
	})
	grayscale := orchestrate.NewGrayscale(rdb, orchLC, nil)

	// ── Phase 1: Blocked-node detection ───────────────────────────────────────
	//
	// Inject: 2/3 domestic ISPs fail, overseas OK → GFW-block pattern.

	blockSnaps := snapsForNode(
		cnFail("ChinaTelecom"), cnFail("ChinaUnicom"), cnOK("ChinaMobile"), overseasOK(),
	)
	snapper.setNode(e2eNode, blockSnaps)

	// 2 detect ticks → up → blocked_suspect, weight → 10.
	for i := 0; i < 2; i++ {
		if err := engine.Tick(ctx); err != nil {
			t.Fatalf("detect tick %d: %v", i, err)
		}
	}
	if got := detectLC.NodeStatus(e2eNode); got != detect.StatusBlockedSuspect {
		t.Fatalf("after 2 ticks: status=%q want blocked_suspect", got)
	}
	if got := detectLC.NodeWeight(e2eNode); got != detectCfg.SuspectWeight {
		t.Errorf("suspect weight=%d want %d", got, detectCfg.SuspectWeight)
	}

	// 6 more detect ticks → blocked_confirmed → down, replace queue populated.
	for i := 0; i < detectCfg.ConfirmedStreakMin; i++ {
		if err := engine.Tick(ctx); err != nil {
			t.Fatalf("confirm tick %d: %v", i, err)
		}
	}
	if got := detectLC.NodeStatus(e2eNode); got != detect.StatusDown {
		t.Fatalf("after 8 ticks: status=%q want down", got)
	}

	// Verify replace queue.
	qlen, _ := rdb.LLen(ctx, "detect:replace:queue").Result()
	if qlen != 1 {
		t.Fatalf("replace queue length=%d want 1", qlen)
	}
	raw, _ := rdb.LIndex(ctx, "detect:replace:queue", 0).Result()
	var qentry struct {
		NodeID          string `json:"nodeId"`
		ReplacementUUID string `json:"replacementUuid"`
	}
	if err := json.Unmarshal([]byte(raw), &qentry); err != nil {
		t.Fatalf("unmarshal queue entry: %v", err)
	}
	if qentry.NodeID != e2eNode {
		t.Errorf("queue nodeId=%q want %q", qentry.NodeID, e2eNode)
	}
	repUUID := qentry.ReplacementUUID

	// ── Phase 2: Replacement orchestration ────────────────────────────────────
	//
	// Orchestrate ticks drive the state machine: pending → creating → probing
	// → activating → draining_old → done.

	// Tick 1: drainQueue (pending) + stepPending (→ creating).
	if err := replacer.Tick(ctx); err != nil {
		t.Fatalf("orch tick 1: %v", err)
	}

	// Tick 2: stepCreating → CreateNode → phase=probing.
	if err := replacer.Tick(ctx); err != nil {
		t.Fatalf("orch tick 2: %v", err)
	}
	if prov.createCount() != 1 {
		t.Fatalf("expected 1 CreateNode call after 2 orch ticks, got %d", prov.createCount())
	}
	newNodeID := prov.lastCreatedID()
	if newNodeID == "" {
		t.Fatal("CreateNode returned empty nodeID")
	}
	t.Logf("replacement node created: %s", newNodeID)

	// Inject healthy probe data for the new node.
	goodSnaps := snapsForNode(
		cnOK("ChinaTelecom"), cnOK("ChinaUnicom"), cnOK("ChinaMobile"), overseasOK(),
	)
	snapper.setNode(newNodeID, goodSnaps)

	// Tick 3: stepProbing → probeStreak=1 (need 2).
	if err := replacer.Tick(ctx); err != nil {
		t.Fatalf("orch tick 3: %v", err)
	}

	// Tick 4: stepProbing → probeStreak=2 ≥ ProbeCyclesRequired → phase=activating.
	if err := replacer.Tick(ctx); err != nil {
		t.Fatalf("orch tick 4: %v", err)
	}

	// Tick 5: stepActivating → SetWeight(10) + probing→up + BumpVersion + startGrayscale.
	if err := replacer.Tick(ctx); err != nil {
		t.Fatalf("orch tick 5: %v", err)
	}

	// Tick 6: stepDrainingOld → DestroyNode(old) + breaker.Record + phase=done + audit.
	if err := replacer.Tick(ctx); err != nil {
		t.Fatalf("orch tick 6: %v", err)
	}

	// ── Phase 3: Assertions ───────────────────────────────────────────────────

	// 1. Detect lifecycle events: up→suspect, suspect→confirmed, confirmed→down.
	events := detectLC.Events()
	wantEvents := [][2]detect.NodeStatus{
		{detect.StatusUp, detect.StatusBlockedSuspect},
		{detect.StatusBlockedSuspect, detect.StatusBlockedConfirmed},
		{detect.StatusBlockedConfirmed, detect.StatusDown},
	}
	if len(events) != len(wantEvents) {
		t.Errorf("detect events count=%d want %d: %v", len(events), len(wantEvents), events)
	} else {
		for i, ev := range events {
			if ev.From != wantEvents[i][0] || ev.To != wantEvents[i][1] {
				t.Errorf("detect event[%d]: %q→%q want %q→%q",
					i, ev.From, ev.To, wantEvents[i][0], wantEvents[i][1])
			}
		}
	}

	// 2. Version was bumped (directory version bump for client refetch).
	if orchLC.version() == 0 {
		t.Error("BumpVersion not called — directory version not bumped")
	}

	// 3. Old node was destroyed.
	if prov.destroyCount() == 0 {
		t.Error("DestroyNode not called for old node")
	}

	// 4. Orchestrate audit log has replacement_done entry.
	if !orchLC.hasAudit("orchestrate|replacement_done|node:" + e2eNode) {
		t.Errorf("missing audit log entry for replacement_done; entries: %v", orchLC.auditEntries())
	}

	// 5. Grayscale key exists for new node (sched:gray:{newNodeID}).
	grayKey := "sched:gray:" + newNodeID
	if exists, _ := rdb.Exists(ctx, grayKey).Result(); exists == 0 {
		t.Errorf("grayscale key %q not created after activation", grayKey)
	}

	// 6. ReplaceRecord is in terminal phase=done.
	recKey := "sched:replace:" + repUUID
	recRaw, err := rdb.Get(ctx, recKey).Result()
	if err != nil {
		t.Fatalf("replace record not found: %v", err)
	}
	var rec struct{ Phase string `json:"phase"` }
	if err := json.Unmarshal([]byte(recRaw), &rec); err != nil {
		t.Fatalf("unmarshal record: %v", err)
	}
	if rec.Phase != "done" {
		t.Errorf("replace record phase=%q want done", rec.Phase)
	}

	// 7. Grayscale.Advance: call it manually with a fast clock to test weight ramp.
	//    (Production ramp interval is 6 h; we call it directly here.)
	_ = grayscale // advance is tested indirectly via the grayKey existence check above.
}

// ─────────────────────────────────────────────────────────────────────────────
// TestSCHED_ENABLED guard: scheduler Config must accept nil Prober/Targets
// ─────────────────────────────────────────────────────────────────────────────

func TestNilProberDoesNotPanic(t *testing.T) {
	rdb, _ := newTestRedis(t)
	sched := scheduler.New(scheduler.Config{
		RDB:                 rdb,
		InstanceID:          "nil-prober-test",
		Engine:              &stubEngine{},
		Replacer:            &stubReplacer{},
		Grayscale:           &stubGrayscale{},
		DetectInterval:      80 * time.Millisecond,
		OrchestrateInterval: 80 * time.Millisecond,
		CapacityInterval:    80 * time.Millisecond,
		LeaseTTL:            200 * time.Millisecond,
		RenewPeriod:         50 * time.Millisecond,
		RetryPeriod:         50 * time.Millisecond,
		TickTimeout:         5 * time.Second,
		// Prober and Targets intentionally nil
	})

	ctx, cancel := context.WithTimeout(context.Background(), 300*time.Millisecond)
	defer cancel()

	// Must not panic.
	if err := sched.Start(ctx); err != nil {
		t.Errorf("Start returned error: %v", err)
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Mock implementations for the e2e scenario
// ─────────────────────────────────────────────────────────────────────────────

// ── mockSnapshotter ───────────────────────────────────────────────────────────
//
// Satisfies both detect.ProbeSnapshotter and orchestrate.ProbeSnapshotter
// (identical interface signatures).

type mockSnapshotter struct {
	mu   sync.Mutex
	data map[string]map[string]probe.ProbeSnapshot
}

func (m *mockSnapshotter) SnapshotsByNode(_ context.Context, nodeID string) (map[string]probe.ProbeSnapshot, error) {
	m.mu.Lock(); defer m.mu.Unlock()
	if snaps, ok := m.data[nodeID]; ok {
		return snaps, nil
	}
	return nil, nil
}

func (m *mockSnapshotter) setNode(nodeID string, snaps map[string]probe.ProbeSnapshot) {
	m.mu.Lock(); defer m.mu.Unlock()
	m.data[nodeID] = snaps
}

// ── mockOrchestrateLC ─────────────────────────────────────────────────────────
//
// Implements orchestrate.LifecycleService for the e2e test.

type mockOrchestrateLC struct {
	mu        sync.Mutex
	nodes     map[string]*orchestrate.NodeInfo
	weights   map[string]int
	ver       int64
	auditLogs []string
	trans     []orchTransEvent
}

type orchTransEvent struct{ NodeID, From, To string }

func newMockOrchestrateLC() *mockOrchestrateLC {
	return &mockOrchestrateLC{
		nodes:   make(map[string]*orchestrate.NodeInfo),
		weights: make(map[string]int),
	}
}

func (m *mockOrchestrateLC) addNode(n *orchestrate.NodeInfo) {
	m.mu.Lock(); defer m.mu.Unlock()
	cp := *n; m.nodes[n.ID] = &cp
}

func (m *mockOrchestrateLC) GetNode(_ context.Context, nodeID string) (*orchestrate.NodeInfo, error) {
	m.mu.Lock(); defer m.mu.Unlock()
	n, ok := m.nodes[nodeID]
	if !ok { return nil, nil }
	cp := *n; return &cp, nil
}

func (m *mockOrchestrateLC) TransitionStatus(_ context.Context, nodeID, from, to string, _ map[string]any) (int, error) {
	m.mu.Lock(); defer m.mu.Unlock()
	m.trans = append(m.trans, orchTransEvent{nodeID, from, to})
	return 1, nil
}

func (m *mockOrchestrateLC) SetWeight(_ context.Context, nodeID string, weight int) error {
	m.mu.Lock(); defer m.mu.Unlock()
	m.weights[nodeID] = weight
	return nil
}

func (m *mockOrchestrateLC) BumpVersion(_ context.Context) error {
	m.mu.Lock(); defer m.mu.Unlock()
	m.ver++; return nil
}

func (m *mockOrchestrateLC) WriteAuditLog(_ context.Context, actor, action, target, _ string) error {
	m.mu.Lock(); defer m.mu.Unlock()
	m.auditLogs = append(m.auditLogs, actor+"|"+action+"|"+target)
	return nil
}

func (m *mockOrchestrateLC) version() int64 {
	m.mu.Lock(); defer m.mu.Unlock(); return m.ver
}

func (m *mockOrchestrateLC) hasAudit(entry string) bool {
	m.mu.Lock(); defer m.mu.Unlock()
	for _, l := range m.auditLogs { if l == entry { return true } }
	return false
}

func (m *mockOrchestrateLC) auditEntries() []string {
	m.mu.Lock(); defer m.mu.Unlock()
	return append([]string{}, m.auditLogs...)
}

// ── mockProvision ─────────────────────────────────────────────────────────────

type mockProvision struct {
	mu           sync.Mutex
	providers    []orchestrate.ProviderInfo
	createCalls  []string // nodeIDs
	destroyCalls []string
	seq          int
	idem         map[string]string
}

func newMockProvision(providers ...orchestrate.ProviderInfo) *mockProvision {
	return &mockProvision{providers: providers, idem: map[string]string{}}
}

func (m *mockProvision) CreateNode(_ context.Context, _ orchestrate.NodeSpec, idemKey string) (string, error) {
	m.mu.Lock(); defer m.mu.Unlock()
	if existing, ok := m.idem[idemKey]; ok { return existing, nil }
	m.seq++
	id := fmt.Sprintf("new-node-%d", m.seq)
	m.idem[idemKey] = id
	m.createCalls = append(m.createCalls, id)
	return id, nil
}

func (m *mockProvision) DestroyNode(_ context.Context, nodeID string) error {
	m.mu.Lock(); defer m.mu.Unlock()
	m.destroyCalls = append(m.destroyCalls, nodeID)
	return nil
}

func (m *mockProvision) RotateIP(_ context.Context, _ string) (string, error) { return "", nil }

func (m *mockProvision) ListProviders(_ context.Context, _, _ string) ([]orchestrate.ProviderInfo, error) {
	m.mu.Lock(); defer m.mu.Unlock()
	return append([]orchestrate.ProviderInfo{}, m.providers...), nil
}

func (m *mockProvision) createCount() int {
	m.mu.Lock(); defer m.mu.Unlock(); return len(m.createCalls)
}

func (m *mockProvision) destroyCount() int {
	m.mu.Lock(); defer m.mu.Unlock(); return len(m.destroyCalls)
}

func (m *mockProvision) lastCreatedID() string {
	m.mu.Lock(); defer m.mu.Unlock()
	if len(m.createCalls) == 0 { return "" }
	return m.createCalls[len(m.createCalls)-1]
}

// ─────────────────────────────────────────────────────────────────────────────
// Probe snapshot builder helpers (mirrors detect/engine_test.go)
// ─────────────────────────────────────────────────────────────────────────────

func snap(country, isp string, l1OK bool, l3OK *bool) probe.ProbeSnapshot {
	r := probe.NodeReport{L1: probe.L1Result{OK: l1OK}}
	if l3OK != nil { r.L3 = &probe.L3Result{OK: *l3OK} }
	return probe.ProbeSnapshot{
		Vantage: probe.VantagePoint{Country: country, ISP: isp},
		Report:  r,
	}
}

func boolPtr(b bool) *bool { return &b }

func cnFail(isp string) probe.ProbeSnapshot { return snap("CN", isp, true, boolPtr(false)) }
func cnOK(isp string) probe.ProbeSnapshot   { return snap("CN", isp, true, boolPtr(true)) }
func overseasOK() probe.ProbeSnapshot       { return snap("SG", "AWS", true, nil) }

func snapsForNode(snaps ...probe.ProbeSnapshot) map[string]probe.ProbeSnapshot {
	m := make(map[string]probe.ProbeSnapshot, len(snaps))
	for _, s := range snaps {
		key := s.Vantage.Country + ":" + s.Vantage.ISP
		m[key] = s
	}
	return m
}