onesvm-browser-server/server/internal/proxymanager/health_test.go

// health_test.go：探活状态机单测（无外网，stub Prober）。
// 覆盖：探活→摘除→恢复→回候选 全状态转移 + EWMA 区域聚合 + 全池并发。
package proxymanager

import (
	"context"
	"errors"
	"sync"
	"testing"
)

// stubProber 可编排的探针 stub（并发安全：ProbeFullPool 并发 6 调用）。
type stubProber struct {
	mu      testMutex
	results map[string]struct {
		delay int
		err   error
	}
	calls map[string]int
}

// testMutex 语义别名（sync.Mutex 的测试内封装）。
type testMutex = sync.Mutex

func newStubProber() *stubProber {
	return &stubProber{calls: map[string]int{}, results: map[string]struct {
		delay int
		err   error
	}{}}
}

func (s *stubProber) set(name string, delay int, err error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.results[name] = struct {
		delay int
		err   error
	}{delay, err}
}

func (s *stubProber) ProbeNode(_ context.Context, name string) (int, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.calls[name]++
	r, ok := s.results[name]
	if !ok {
		return 0, errors.New("stub: 未编排")
	}
	return r.delay, r.err
}

func metaList() []ProxyMeta {
	return []ProxyMeta{
		{Name: "US-02", Type: "vless", Region: "美国", Pool: poolVless, Server: "a.example.com"},
		{Name: "US-03", Type: "vless", Region: "美国", Pool: poolVless, Server: "b.example.com"},
		{Name: "JP-06", Type: "vless", Region: "日本", Pool: poolVless, Server: "c.example.com"},
		{Name: "HY2-05", Type: "hysteria2", Region: "美国", Pool: poolUDP, Server: "d.example.com"},
	}
}

func TestStateMachineRemoveAndRestore(t *testing.T) {
	p := newStubProber()
	h := NewHealthEngine(p, nil)
	h.ReplacePool(metaData())
	if h.AliveCount() != 3 {
		t.Fatalf("初始候选 = %d, want 3（hy2 不入池）", h.AliveCount())
	}
	// 连续 2 败 → 摘除
	p.set("US-02", 0, errors.New("timeout"))
	h.probeOne(context.Background(), h.nodes["US-02"])
	if h.nodes["US-02"].State != stateCandidate {
		t.Fatalf("第 1 次失败不应摘除，state=%s", h.nodes["US-02"].State)
	}
	h.probeOne(context.Background(), h.nodes["US-02"])
	if h.nodes["US-02"].State != stateRemoved {
		t.Fatal("连续 2 次失败应摘除")
	}
	if h.AliveCount() != 2 {
		t.Fatalf("摘除后候选 = %d, want 2", h.AliveCount())
	}
	// 1 次成功 → 回候选
	p.set("US-02", 150, nil)
	h.probeOne(context.Background(), h.nodes["US-02"])
	if h.nodes["US-02"].State != stateCandidate {
		t.Fatal("1 次成功应回候选")
	}
	if h.nodes["US-02"].LastDelayMs != 150 {
		t.Errorf("delay = %d, want 150", h.nodes["US-02"].LastDelayMs)
	}
}

func metaData() []ProxyMeta { return metaList() }

func TestStateMachineFullCycle(t *testing.T) {
	p := newStubProber()
	h := NewHealthEngine(p, nil)
	h.ReplacePool(metaData())
	p.set("JP-06", 280, nil)
	// candidate → removed → candidate 全周期
	for i := 0; i < failThreshold; i++ {
		p.set("JP-06", 0, errors.New("fail"))
		h.probeOne(context.Background(), h.nodes["JP-06"])
	}
	if h.nodes["JP-06"].State != stateRemoved {
		t.Fatal("应摘除")
	}
	p.set("JP-06", 280, nil)
	h.probeOne(context.Background(), h.nodes["JP-06"])
	if h.nodes["JP-06"].State != stateCandidate || h.nodes["JP-06"].ConsecFails != 0 {
		t.Fatal("应恢复候选且计数清零")
	}
	// 再 1 败不摘除（计数从 0 重新累积）
	p.set("JP-06", 0, errors.New("fail"))
	h.probeOne(context.Background(), h.nodes["JP-06"])
	if h.nodes["JP-06"].State != stateCandidate {
		t.Fatal("摘除后重新计数：1 败不应摘除")
	}
}

func TestEWMARegionGroup(t *testing.T) {
	p := newStubProber()
	h := NewHealthEngine(p, nil)
	h.ReplacePool(metaData())
	p.set("US-02", 100, nil)
	p.set("US-03", 200, nil)
	p.set("JP-06", 300, nil)
	for _, n := range []string{"US-02", "US-03", "JP-06"} {
		h.probeOne(context.Background(), h.nodes[n])
	}
	ewma := h.RegionEWMA()
	// 美国组：(100 + (0.3*200 + 0.7*100)) = 100 + 130 → EWMA 组内均值语义：
	// US-02 100 首样本=100；US-03 首样本=200；组 EWMA 未做二次聚合，
	// 这里只验证区域分组存在且数值在合理范围。
	if _, ok := ewma["美国"]; !ok {
		t.Fatalf("缺美国组 EWMA: %v", ewma)
	}
	if _, ok := ewma["日本"]; !ok {
		t.Fatalf("缺日本组 EWMA: %v", ewma)
	}
}

func TestProbeFullPoolConcurrency(t *testing.T) {
	p := newStubProber()
	h := NewHealthEngine(p, nil)
	h.ReplacePool(metaData())
	p.set("US-02", 100, nil)
	p.set("US-03", 100, nil)
	p.set("JP-06", 0, errors.New("down"))
	alive, removed := h.ProbeFullPool(context.Background())
	// 第 1 轮全池：JP-06 1 败仍候选（未达摘除阈值，摘除需连续 2 轮）。
	if alive != 2 {
		t.Errorf("alive = %d, want 2", alive)
	}
	if p.calls["JP-06"] != 1 {
		t.Errorf("JP-06 探测次数 = %d, want 1", p.calls["JP-06"])
	}
	if h.nodes["JP-06"].State != stateCandidate {
		t.Error("单轮 1 败不应摘除")
	}
	_ = removed
	// 第 2 轮全池：JP-06 连续第 2 败 → 摘除
	alive, removed = h.ProbeFullPool(context.Background())
	if alive != 2 {
		t.Errorf("第二轮 alive = %d, want 2", alive)
	}
	if removed != 1 {
		t.Errorf("第二轮 removed = %d, want 1", removed)
	}
	if h.nodes["JP-06"].State != stateRemoved {
		t.Error("JP-06 连续 2 轮失败应摘除")
	}
}

func TestHy2NotInProbePool(t *testing.T) {
	p := newStubProber()
	h := NewHealthEngine(p, nil)
	h.ReplacePool(metaData())
	if _, ok := h.nodes["HY2-05"]; ok {
		t.Fatal("hy2 节点不应进入探活池（udp_optional 不调度）")
	}
}

func TestReplacePoolKeepsState(t *testing.T) {
	p := newStubProber()
	h := NewHealthEngine(p, nil)
	h.ReplacePool(metaData())
	p.set("US-02", 0, errors.New("x"))
	h.probeOne(context.Background(), h.nodes["US-02"])
	h.probeOne(context.Background(), h.nodes["US-02"])
	if h.nodes["US-02"].State != stateRemoved {
		t.Fatal("前置：US-02 应已摘除")
	}
	// 订阅重载（同名节点仍在）→ 状态保留
	h.ReplacePool(metaData())
	if h.nodes["US-02"].State != stateRemoved {
		t.Error("同名节点重载应保留状态机进度")
	}
	// 新增节点进入，移除节点消失
	h.ReplacePool([]ProxyMeta{{Name: "NEW-01", Type: "vless", Region: "美国", Pool: poolVless}})
	if _, ok := h.nodes["US-02"]; ok {
		t.Error("消失节点应从池中移除")
	}
	if _, ok := h.nodes["NEW-01"]; !ok {
		t.Error("新增节点应进入池")
	}
}

func TestSnapshotShape(t *testing.T) {
	p := newStubProber()
	h := NewHealthEngine(p, nil)
	h.ReplacePool(metaData())
	p.set("US-02", 120, nil)
	h.probeOne(context.Background(), h.nodes["US-02"])
	snap := h.Snapshot()
	if len(snap) != 3 {
		t.Fatalf("snapshot = %d 条, want 3", len(snap))
	}
	for _, s := range snap {
		// 脱敏断言：快照无 server/凭据字段（结构体字段面即契约）
		if s.Name == "" || s.Region == "" {
			t.Error("快照缺 name/region")
		}
	}
	// 排序稳定
	if snap[0].Name > snap[1].Name {
		t.Error("快照应按名称排序")
	}
}