pangolin/server/internal/scheduler/orchestrate/breaker_test.go

package orchestrate_test

import (
	"context"
	"sync"
	"testing"
	"time"

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

// ─────────────────────────────────────────────────────────────────────────────
// Mock PoolReader
// ─────────────────────────────────────────────────────────────────────────────

// mockPoolReader is a simple in-memory PoolReader for tests.
type mockPoolReader struct {
	mu    sync.Mutex
	pools map[string]orchestrate.PoolStat // key = tier+":"+region
}

func newMockPoolReader() *mockPoolReader {
	return &mockPoolReader{pools: make(map[string]orchestrate.PoolStat)}
}

func (r *mockPoolReader) setPool(tier, region string, up, target int) {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.pools[tier+":"+region] = orchestrate.PoolStat{
		Tier: tier, Region: region, UpCount: up, Target: target,
	}
}

func (r *mockPoolReader) AllPools(_ context.Context) ([]orchestrate.PoolStat, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := make([]orchestrate.PoolStat, 0, len(r.pools))
	for _, s := range r.pools {
		out = append(out, s)
	}
	return out, nil
}

func (r *mockPoolReader) PoolCapacity(_ context.Context, tier, region string) (int, int, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	s, ok := r.pools[tier+":"+region]
	if !ok {
		return 0, 0, nil
	}
	return s.UpCount, s.Target, nil
}

// ─────────────────────────────────────────────────────────────────────────────
// Helper: build a RedisBreaker for tests
// ─────────────────────────────────────────────────────────────────────────────

func newTestBreaker(t *testing.T, pools *mockPoolReader, notifier *mockNotifier, clock *mockClock) (*orchestrate.RedisBreaker, *orchestrate.ConfigManager) {
	t.Helper()
	rdb, _ := newTestRedis(t)
	cfgMgr := orchestrate.NewConfigManager("", nil)
	_ = cfgMgr.Load()

	return orchestrate.NewRedisBreaker(orchestrate.RedisBreakerConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		CfgMgr:   cfgMgr,
		Clock:    clock,
	}), cfgMgr
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: N-th Record trips the breaker; (N+1)-th Allow returns false
// ─────────────────────────────────────────────────────────────────────────────

func TestBreakerTripsAtThreshold(t *testing.T) {
	const (
		tier   = "free"
		region = "hkg"
	)

	pools := newMockPoolReader()
	// Pool of 10 nodes → threshold = ceil(10 × 30%) = 3.
	pools.setPool(tier, region, 10, 10)

	notifier := &mockNotifier{}
	clock := newMockClock(time.Unix(1_700_000_000, 0).UTC())

	b, _ := newTestBreaker(t, pools, notifier, clock)

	// Before any records: Allow should return true.
	if !b.Allow(tier, region) {
		t.Fatal("Allow should return true before any records")
	}

	// Record N-1 replacements (threshold-1 = 2); Allow should still be true.
	b.Record(tier, region) // count = 1
	b.Record(tier, region) // count = 2
	if !b.Allow(tier, region) {
		t.Fatalf("Allow should return true with count < threshold (count=2, threshold=3)")
	}

	// Record the N-th replacement; now count = 3 = threshold.
	b.Record(tier, region) // count = 3

	// (N+1)-th Allow: breaker is tripped → false.
	if b.Allow(tier, region) {
		t.Error("Allow should return false when count >= threshold (tripped)")
	}

	// Critical alert must have been emitted exactly once.
	if n := notifier.count(); n != 1 {
		t.Errorf("NotifyFault calls = %d; want 1 (on trip)", n)
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Window slide-out auto-recovers the breaker
// ─────────────────────────────────────────────────────────────────────────────

func TestBreakerAutoRecoveryAfterWindowExpiry(t *testing.T) {
	const (
		tier   = "premium"
		region = "sin"
	)

	pools := newMockPoolReader()
	// Pool of 10 → threshold = 3.
	pools.setPool(tier, region, 10, 10)

	notifier := &mockNotifier{}
	clock := newMockClock(time.Unix(1_700_000_000, 0).UTC())
	b, _ := newTestBreaker(t, pools, notifier, clock)

	// Record N entries to trip.
	b.Record(tier, region)
	b.Record(tier, region)
	b.Record(tier, region)
	if b.Allow(tier, region) {
		t.Fatal("breaker should be tripped after 3 records (threshold=3)")
	}

	// Advance clock past the 1-h window; all ZSET entries are now stale.
	clock.advance(61 * time.Minute)

	// Allow should return true: ZREMRANGEBYSCORE removes stale entries,
	// count drops to 0 < threshold, and the trip flag TTL has expired.
	if !b.Allow(tier, region) {
		t.Error("Allow should return true after the window has expired (auto-recovery)")
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Reset manually clears the breaker with audit
// ─────────────────────────────────────────────────────────────────────────────

func TestBreakerResetRestoresAllow(t *testing.T) {
	const (
		tier   = "free"
		region = "tyo"
	)

	pools := newMockPoolReader()
	pools.setPool(tier, region, 10, 10) // threshold = 3

	notifier := &mockNotifier{}
	lc := newMockLC()
	clock := newMockClock(time.Unix(1_700_000_000, 0).UTC())
	rdb, _ := newTestRedis(t)
	cfgMgr := orchestrate.NewConfigManager("", nil)
	_ = cfgMgr.Load()

	b := orchestrate.NewRedisBreaker(orchestrate.RedisBreakerConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		LC:       lc,
		CfgMgr:   cfgMgr,
		Clock:    clock,
	})

	// Trip the breaker.
	b.Record(tier, region)
	b.Record(tier, region)
	b.Record(tier, region)
	if b.Allow(tier, region) {
		t.Fatal("breaker should be tripped")
	}

	// Reset.
	ctx := context.Background()
	if err := b.Reset(ctx, tier, region, "admin-alice"); err != nil {
		t.Fatalf("Reset error: %v", err)
	}

	// Allow should now return true.
	if !b.Allow(tier, region) {
		t.Error("Allow should return true after Reset")
	}

	// Audit log must mention the reset.
	if !lc.hasAudit("admin-alice|breaker_reset|pool:" + tier + ":" + region) {
		t.Errorf("audit log missing breaker_reset entry; got: %v", lc.auditLogs)
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: BreakerMinN = 3 enforced on tiny pool (target < 10)
// ─────────────────────────────────────────────────────────────────────────────

func TestBreakerMinNEnforcedOnSmallPool(t *testing.T) {
	const (
		tier   = "free"
		region = "fra"
	)

	pools := newMockPoolReader()
	// Pool of 2 nodes; 30% of 2 = 0.6 → ceil = 1, but min = 3.
	pools.setPool(tier, region, 2, 2)

	notifier := &mockNotifier{}
	clock := newMockClock(time.Unix(1_700_000_000, 0).UTC())
	b, _ := newTestBreaker(t, pools, notifier, clock)

	// 2 records should NOT trip (threshold = 3, not 1).
	b.Record(tier, region)
	b.Record(tier, region)
	if !b.Allow(tier, region) {
		t.Error("Allow should return true: count=2, threshold=min(3)")
	}

	// 3rd record trips.
	b.Record(tier, region)
	if b.Allow(tier, region) {
		t.Error("Allow should return false: count=3 >= minN(3)")
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Trip alert is emitted only once per window epoch (dedup)
// ─────────────────────────────────────────────────────────────────────────────

func TestBreakerAlertSentOnlyOnce(t *testing.T) {
	const (
		tier   = "free"
		region = "ams"
	)

	pools := newMockPoolReader()
	pools.setPool(tier, region, 10, 10) // threshold = 3

	notifier := &mockNotifier{}
	clock := newMockClock(time.Unix(1_700_000_000, 0).UTC())
	b, _ := newTestBreaker(t, pools, notifier, clock)

	// Trip.
	b.Record(tier, region)
	b.Record(tier, region)
	b.Record(tier, region)

	// Call Allow multiple times while tripped.
	for i := 0; i < 5; i++ {
		b.Allow(tier, region)
	}

	// NotifyFault must have been called exactly once (from Record reaching N).
	if n := notifier.count(); n != 1 {
		t.Errorf("NotifyFault calls = %d; want exactly 1 (trip alert dedup)", n)
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: pool not found → falls back to BreakerMinN (3)
// ─────────────────────────────────────────────────────────────────────────────

func TestBreakerFallsBackToMinNForUnknownPool(t *testing.T) {
	pools := newMockPoolReader()
	// No pool registered for "unknown"/"xyz".

	notifier := &mockNotifier{}
	clock := newMockClock(time.Unix(1_700_000_000, 0).UTC())
	b, _ := newTestBreaker(t, pools, notifier, clock)

	b.Record("unknown", "xyz")
	b.Record("unknown", "xyz")
	if !b.Allow("unknown", "xyz") {
		t.Error("Allow should return true: count=2, threshold=minN(3) for unknown pool")
	}

	b.Record("unknown", "xyz")
	if b.Allow("unknown", "xyz") {
		t.Error("Allow should return false: count=3 >= minN(3)")
	}
}