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

package orchestrate_test

import (
	"context"
	"os"
	"testing"
	"time"

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

// ─────────────────────────────────────────────────────────────────────────────
// Utility: write a temporary YAML config file
// ─────────────────────────────────────────────────────────────────────────────

func writeTempYAML(t *testing.T, content string) (string, error) {
	t.Helper()
	f, err := os.CreateTemp("", "sched_config_*.yaml")
	if err != nil {
		return "", err
	}
	t.Cleanup(func() { os.Remove(f.Name()) })
	if _, err := f.WriteString(content); err != nil {
		f.Close()
		return "", err
	}
	f.Close()
	return f.Name(), nil
}

// defaultCfgMgr returns a ConfigManager loaded with production defaults.
func defaultCfgMgr(t *testing.T) *orchestrate.ConfigManager {
	t.Helper()
	m := orchestrate.NewConfigManager("", nil)
	if err := m.Load(); err != nil {
		t.Fatalf("ConfigManager.Load: %v", err)
	}
	return m
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Pool below 70% watermark → alert emitted
// ─────────────────────────────────────────────────────────────────────────────

func TestCapacityWatermarkAlert(t *testing.T) {
	rdb, _ := newTestRedis(t)
	pools := newMockPoolReader()
	// 6 out of 10 up = 60% < 70% → alert expected.
	pools.setPool("free", "hkg", 6, 10)

	notifier := &mockNotifier{}
	cm := orchestrate.NewCapacityMonitor(orchestrate.CapacityMonitorConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		CfgMgr:   defaultCfgMgr(t),
	})

	if err := cm.RunOnce(context.Background()); err != nil {
		t.Fatalf("RunOnce: %v", err)
	}

	if notifier.count() != 1 {
		t.Errorf("NotifyFault calls = %d; want 1 (watermark alert)", notifier.count())
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Pool at exactly 70% → no alert
// ─────────────────────────────────────────────────────────────────────────────

func TestCapacityWatermarkNoAlertAtThreshold(t *testing.T) {
	rdb, _ := newTestRedis(t)
	pools := newMockPoolReader()
	// 7/10 = 70% == threshold → healthy, no alert.
	pools.setPool("free", "hkg", 7, 10)

	notifier := &mockNotifier{}
	cm := orchestrate.NewCapacityMonitor(orchestrate.CapacityMonitorConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		CfgMgr:   defaultCfgMgr(t),
	})

	if err := cm.RunOnce(context.Background()); err != nil {
		t.Fatalf("RunOnce: %v", err)
	}

	if notifier.count() != 0 {
		t.Errorf("NotifyFault calls = %d; want 0 (fill rate = 70%% >= threshold)", notifier.count())
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Repeated RunOnce within suppress window → only 1 alert
// ─────────────────────────────────────────────────────────────────────────────

func TestCapacityWatermarkAlertSuppression(t *testing.T) {
	rdb, mr := newTestRedis(t)
	pools := newMockPoolReader()
	pools.setPool("free", "sin", 5, 10) // 50% < 70%

	notifier := &mockNotifier{}
	cm := orchestrate.NewCapacityMonitor(orchestrate.CapacityMonitorConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		CfgMgr:   defaultCfgMgr(t),
	})

	ctx := context.Background()

	// Three RunOnce calls within the 10-min suppress window.
	for i := 0; i < 3; i++ {
		if err := cm.RunOnce(ctx); err != nil {
			t.Fatalf("RunOnce %d: %v", i, err)
		}
	}
	if notifier.count() != 1 {
		t.Errorf("NotifyFault calls = %d; want 1 (suppressed repeats)", notifier.count())
	}

	// Advance miniredis TTL past the 10-min suppress window.
	mr.FastForward(11 * time.Minute)

	// After suppress window expires → another alert.
	if err := cm.RunOnce(ctx); err != nil {
		t.Fatalf("RunOnce after expire: %v", err)
	}
	if notifier.count() != 2 {
		t.Errorf("NotifyFault calls = %d; want 2 after suppress window expires", notifier.count())
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Missing probe heartbeat → alert; present probe → no alert
// ─────────────────────────────────────────────────────────────────────────────

func TestCapacityProbeHeartbeatMissingAlert(t *testing.T) {
	rdb, _ := newTestRedis(t)
	pools := newMockPoolReader()
	notifier := &mockNotifier{}

	ctx := context.Background()
	// probe-b has a live heartbeat; probe-a does not.
	if err := rdb.Set(ctx, "probe:hb:probe-b", "1", 15*time.Minute).Err(); err != nil {
		t.Fatalf("seed probe-b heartbeat: %v", err)
	}

	yamlContent := `
capacity:
  expected_probe_ids:
    - probe-a
    - probe-b
  alert_suppress_min: 10
`
	tmpFile, err := writeTempYAML(t, yamlContent)
	if err != nil {
		t.Fatalf("write temp YAML: %v", err)
	}
	cfgMgr := orchestrate.NewConfigManager(tmpFile, nil)
	if err := cfgMgr.Load(); err != nil {
		t.Fatalf("Load YAML config: %v", err)
	}

	cm := orchestrate.NewCapacityMonitor(orchestrate.CapacityMonitorConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		CfgMgr:   cfgMgr,
	})

	if err := cm.RunOnce(ctx); err != nil {
		t.Fatalf("RunOnce: %v", err)
	}

	// Only probe-a should trigger an alert.
	if notifier.count() != 1 {
		t.Errorf("NotifyFault calls = %d; want 1 (only probe-a missing)", notifier.count())
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Both probes alive → no alert
// ─────────────────────────────────────────────────────────────────────────────

func TestCapacityProbeHeartbeatBothAlive(t *testing.T) {
	rdb, _ := newTestRedis(t)
	pools := newMockPoolReader()
	notifier := &mockNotifier{}

	ctx := context.Background()
	rdb.Set(ctx, "probe:hb:probe-a", "1", 15*time.Minute)
	rdb.Set(ctx, "probe:hb:probe-b", "1", 15*time.Minute)

	yamlContent := `
capacity:
  expected_probe_ids:
    - probe-a
    - probe-b
`
	tmpFile, err := writeTempYAML(t, yamlContent)
	if err != nil {
		t.Fatalf("write temp YAML: %v", err)
	}
	cfgMgr := orchestrate.NewConfigManager(tmpFile, nil)
	if err := cfgMgr.Load(); err != nil {
		t.Fatalf("Load YAML: %v", err)
	}

	cm := orchestrate.NewCapacityMonitor(orchestrate.CapacityMonitorConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		CfgMgr:   cfgMgr,
	})

	if err := cm.RunOnce(ctx); err != nil {
		t.Fatalf("RunOnce: %v", err)
	}

	if notifier.count() != 0 {
		t.Errorf("NotifyFault calls = %d; want 0 (both probes alive)", notifier.count())
	}
}

// ─────────────────────────────────────────────────────────────────────────────
// Test: Multiple pools — only below-threshold pools alert
// ─────────────────────────────────────────────────────────────────────────────

func TestCapacityMultiplePoolsSelectiveAlert(t *testing.T) {
	rdb, _ := newTestRedis(t)
	pools := newMockPoolReader()
	// hkg: 60% → alert
	pools.setPool("free", "hkg", 6, 10)
	// sin: 90% → no alert
	pools.setPool("free", "sin", 9, 10)
	// tyo: 50% → alert
	pools.setPool("free", "tyo", 5, 10)

	notifier := &mockNotifier{}
	cm := orchestrate.NewCapacityMonitor(orchestrate.CapacityMonitorConfig{
		RDB:      rdb,
		Pools:    pools,
		Notifier: notifier,
		CfgMgr:   defaultCfgMgr(t),
	})

	if err := cm.RunOnce(context.Background()); err != nil {
		t.Fatalf("RunOnce: %v", err)
	}

	// Expect 2 alerts: hkg and tyo; sin is healthy.
	if n := notifier.count(); n != 2 {
		t.Errorf("NotifyFault calls = %d; want 2 (hkg + tyo)", n)
	}
}