fix(selector): safe type assertion, clean up debug logging, add doc comments and 64 tests

Replace bare type assertion in ParallelStrategy with comma-ok check to
prevent panic on non-Node inputs. Downgrade noisy Infof hash-selection
log to Tracef and remove Chinese-language debug output. Add doc comments
to all exported symbols in weighted.go.
This commit is contained in:
ginuerzh
2026-05-25 23:38:35 +08:00
parent 18f39f940b
commit aabebd047b
9 changed files with 992 additions and 5 deletions
+266
View File
@@ -0,0 +1,266 @@
package selector
import (
"context"
"testing"
"time"
"github.com/go-gost/core/metadata"
"github.com/go-gost/core/selector"
xmd "github.com/go-gost/x/metadata"
mdutil "github.com/go-gost/x/metadata/util"
)
// testMarkable is a test item that implements both Metadatable and Markable.
type testMarkable struct {
md metadata.Metadata
marker selector.Marker
}
func newTestMarkable(md map[string]any) *testMarkable {
return &testMarkable{
md: xmd.NewMetadata(md),
marker: selector.NewFailMarker(),
}
}
func (t *testMarkable) Metadata() metadata.Metadata {
return t.md
}
func (t *testMarkable) Marker() selector.Marker {
return t.marker
}
// testMetadatable only implements Metadatable, not Markable.
type testMetadatable struct {
md metadata.Metadata
}
func newTestMetadatable(md map[string]any) *testMetadatable {
return &testMetadatable{md: xmd.NewMetadata(md)}
}
func (t *testMetadatable) Metadata() metadata.Metadata {
return t.md
}
// --- FailFilter tests ---
func TestFailFilter_Empty(t *testing.T) {
f := FailFilter[int](1, time.Second)
if vs := f.Filter(context.Background()); len(vs) != 0 {
t.Fatalf("expected empty, got %d items", len(vs))
}
}
func TestFailFilter_Single(t *testing.T) {
f := FailFilter[int](1, time.Second)
vs := f.Filter(context.Background(), 42)
if len(vs) != 1 || vs[0] != 42 {
t.Fatalf("expected [42], got %v", vs)
}
}
func TestFailFilter_HealthyItems(t *testing.T) {
f := FailFilter[*testMarkable](2, 10*time.Second)
items := []*testMarkable{
newTestMarkable(nil),
newTestMarkable(nil),
}
vs := f.Filter(context.Background(), items[0], items[1])
if len(vs) != 2 {
t.Fatalf("expected 2 items, got %d", len(vs))
}
}
func TestFailFilter_DeadItem(t *testing.T) {
f := FailFilter[*testMarkable](1, 10*time.Second)
alive := newTestMarkable(nil)
dead := newTestMarkable(nil)
dead.Marker().Mark() // count=1, within failTimeout
vs := f.Filter(context.Background(), alive, dead)
if len(vs) != 1 || vs[0] != alive {
t.Fatalf("expected only alive item, got %v", vs)
}
}
func TestFailFilter_ExpiredFailTimeout(t *testing.T) {
f := FailFilter[*testMarkable](1, 50*time.Millisecond)
item := newTestMarkable(nil)
item.Marker().Mark()
// Wait for fail timeout to expire
time.Sleep(80 * time.Millisecond)
vs := f.Filter(context.Background(), item)
if len(vs) != 1 {
t.Fatalf("expected item to be included after timeout, got %d items", len(vs))
}
}
func TestFailFilter_PerItemOverrides(t *testing.T) {
f := FailFilter[*testMarkable](1, 10*time.Second)
// Item with custom maxFails=5
item := newTestMarkable(map[string]any{
"maxFails": 5,
"failTimeout": "1h",
})
item.Marker().Mark() // count=1, but maxFails=5 so still alive
vs := f.Filter(context.Background(), item)
if len(vs) != 1 {
t.Fatalf("expected item with high maxFails to pass, got %d items", len(vs))
}
}
func TestFailFilter_NilMarker(t *testing.T) {
f := FailFilter[*testMetadatable](1, time.Second)
item := newTestMetadatable(nil)
vs := f.Filter(context.Background(), item)
if len(vs) != 1 {
t.Fatalf("expected non-markable item to pass, got %d items", len(vs))
}
}
func TestFailFilter_AllDead(t *testing.T) {
f := FailFilter[*testMarkable](1, 10*time.Second)
dead1 := newTestMarkable(nil)
dead1.Marker().Mark()
dead2 := newTestMarkable(nil)
dead2.Marker().Mark()
vs := f.Filter(context.Background(), dead1, dead2)
if len(vs) != 0 {
t.Fatalf("expected all dead items filtered, got %d items", len(vs))
}
}
func TestFailFilter_ZeroMaxFails(t *testing.T) {
// Zero maxFails should default to 1, so a single mark should filter the item
f := FailFilter[*testMarkable](0, 10*time.Second)
item := newTestMarkable(nil)
item.Marker().Mark() // count=1, which >= maxFails(defaulted to 1)
vs := f.Filter(context.Background(), item, newTestMarkable(nil))
if len(vs) != 1 {
t.Fatalf("expected zero maxFails to default to 1, got %d items", len(vs))
}
}
func TestFailFilter_ZeroFailTimeout(t *testing.T) {
// Zero failTimeout should default to DefaultFailTimeout (10s)
f := FailFilter[*testMarkable](1, 0)
item := newTestMarkable(nil)
item.Marker().Mark() // count=1, within DefaultFailTimeout → dead
vs := f.Filter(context.Background(), item, newTestMarkable(nil))
if len(vs) != 1 {
t.Fatalf("expected zero failTimeout to use default, got %d items", len(vs))
}
}
// --- BackupFilter tests ---
func TestBackupFilter_Empty(t *testing.T) {
f := BackupFilter[int]()
if vs := f.Filter(context.Background()); len(vs) != 0 {
t.Fatalf("expected empty, got %d items", len(vs))
}
}
func TestBackupFilter_Single(t *testing.T) {
f := BackupFilter[int]()
vs := f.Filter(context.Background(), 42)
if len(vs) != 1 || vs[0] != 42 {
t.Fatalf("expected [42], got %v", vs)
}
}
func TestBackupFilter_PrimaryPreferred(t *testing.T) {
f := BackupFilter[*testMetadatable]()
primary := newTestMetadatable(nil)
backup := newTestMetadatable(map[string]any{"backup": true})
vs := f.Filter(context.Background(), primary, backup)
if len(vs) != 1 || vs[0] != primary {
t.Fatalf("expected only primary, got %v", vs)
}
}
func TestBackupFilter_AllBackup(t *testing.T) {
f := BackupFilter[*testMetadatable]()
b1 := newTestMetadatable(map[string]any{"backup": true})
b2 := newTestMetadatable(map[string]any{"backup": true})
vs := f.Filter(context.Background(), b1, b2)
if len(vs) != 2 {
t.Fatalf("expected all backups returned when no primaries, got %d items", len(vs))
}
}
func TestBackupFilter_NoBackup(t *testing.T) {
f := BackupFilter[*testMetadatable]()
p1 := newTestMetadatable(nil)
p2 := newTestMetadatable(nil)
vs := f.Filter(context.Background(), p1, p2)
if len(vs) != 2 {
t.Fatalf("expected all primaries, got %d items", len(vs))
}
}
func TestBackupFilter_BackupFalse(t *testing.T) {
f := BackupFilter[*testMetadatable]()
item := newTestMetadatable(map[string]any{"backup": false})
vs := f.Filter(context.Background(), item)
if len(vs) != 1 {
t.Fatalf("expected item with backup=false to be treated as primary, got %d items", len(vs))
}
}
// Verify the constants are accessible
func TestFilterConstants(t *testing.T) {
if DefaultMaxFails != 1 {
t.Fatalf("DefaultMaxFails expected 1, got %d", DefaultMaxFails)
}
if DefaultFailTimeout != 10*time.Second {
t.Fatalf("DefaultFailTimeout expected 10s, got %v", DefaultFailTimeout)
}
}
// Verify metadata label constants via mdutil
func TestFilterMetadataLabels(t *testing.T) {
md := xmd.NewMetadata(map[string]any{
"weight": 5,
"backup": true,
"maxFails": 3,
"failTimeout": "30s",
})
if w := mdutil.GetInt(md, labelWeight); w != 5 {
t.Fatalf("weight: expected 5, got %d", w)
}
if !mdutil.GetBool(md, labelBackup) {
t.Fatal("backup: expected true")
}
if mf := mdutil.GetInt(md, labelMaxFails); mf != 3 {
t.Fatalf("maxFails: expected 3, got %d", mf)
}
if ft := mdutil.GetDuration(md, labelFailTimeout); ft != 30*time.Second {
t.Fatalf("failTimeout: expected 30s, got %v", ft)
}
}
+8 -1
View File
@@ -15,6 +15,8 @@ import (
type parallelStrategy[T any] struct{} type parallelStrategy[T any] struct{}
// ParallelStrategy is a strategy for node selector.
// It dials all nodes concurrently and uses the first successful connection.
func ParallelStrategy[T any]() selector.Strategy[T] { func ParallelStrategy[T any]() selector.Strategy[T] {
return &parallelStrategy[T]{} return &parallelStrategy[T]{}
} }
@@ -29,7 +31,11 @@ func (s *parallelStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
nodes := make([]*chain.Node, 0, len(vs)) nodes := make([]*chain.Node, 0, len(vs))
for _, node := range vs { for _, node := range vs {
nodes = append(nodes, any(node).(*chain.Node)) n, ok := any(node).(*chain.Node)
if !ok {
return
}
nodes = append(nodes, n)
} }
vn := chain.NewNode("parallel", "", chain.TransportNodeOption(&parallelTransporter{nodes: nodes})) vn := chain.NewNode("parallel", "", chain.TransportNodeOption(&parallelTransporter{nodes: nodes}))
@@ -160,6 +166,7 @@ func (tr *parallelTransporter) Options() *chain.TransportOptions {
return &chain.TransportOptions{} return &chain.TransportOptions{}
} }
// Copy implements chain.Transporter.Copy.
func (tr *parallelTransporter) Copy() chain.Transporter { func (tr *parallelTransporter) Copy() chain.Transporter {
return &parallelTransporter{ return &parallelTransporter{
nodes: append([]*chain.Node(nil), tr.nodes...), nodes: append([]*chain.Node(nil), tr.nodes...),
+238
View File
@@ -0,0 +1,238 @@
package selector
import (
"context"
"net"
"testing"
"github.com/go-gost/core/chain"
"github.com/go-gost/core/connector"
)
// mockTransport implements chain.Transporter for testing.
type mockTransport struct {
dialResult net.Conn
dialErr error
handshakeErr error
connectErr error
bindErr error
}
func (m *mockTransport) Dial(ctx context.Context, addr string) (net.Conn, error) {
if m.dialErr != nil {
return nil, m.dialErr
}
return m.dialResult, nil
}
func (m *mockTransport) Handshake(ctx context.Context, conn net.Conn) (net.Conn, error) {
if m.handshakeErr != nil {
return nil, m.handshakeErr
}
return conn, nil
}
func (m *mockTransport) Connect(ctx context.Context, conn net.Conn, network, address string) (net.Conn, error) {
if m.connectErr != nil {
return nil, m.connectErr
}
return conn, nil
}
func (m *mockTransport) Bind(ctx context.Context, conn net.Conn, network, address string, opts ...connector.BindOption) (net.Listener, error) {
if m.bindErr != nil {
return nil, m.bindErr
}
return net.Listen("tcp", "127.0.0.1:0")
}
func (m *mockTransport) Multiplex() bool { return false }
func (m *mockTransport) Options() *chain.TransportOptions { return &chain.TransportOptions{} }
func (m *mockTransport) Copy() chain.Transporter { return m }
func makeNode(tr chain.Transporter) *chain.Node {
return chain.NewNode("test", "127.0.0.1:0", chain.TransportNodeOption(tr))
}
// --- ParallelStrategy ---
func TestParallelStrategy_Empty(t *testing.T) {
s := ParallelStrategy[*chain.Node]()
if v := s.Apply(context.Background()); v != nil {
t.Fatalf("expected nil, got %v", v)
}
}
func TestParallelStrategy_Single(t *testing.T) {
s := ParallelStrategy[*chain.Node]()
tr := &mockTransport{}
node := makeNode(tr)
result := s.Apply(context.Background(), node)
if result == nil {
t.Fatal("expected non-nil result for single node")
}
}
func TestParallelStrategy_MultipleNodes(t *testing.T) {
s := ParallelStrategy[*chain.Node]()
nodes := []*chain.Node{
makeNode(&mockTransport{}),
makeNode(&mockTransport{}),
makeNode(&mockTransport{}),
}
result := s.Apply(context.Background(), nodes[0], nodes[1], nodes[2])
if result == nil {
t.Fatal("expected non-nil result")
}
}
func TestParallelStrategy_NonNodeInput(t *testing.T) {
s := ParallelStrategy[string]()
// Single non-node input: len==1 returns vs[0] directly
if v := s.Apply(context.Background(), "hello"); v != "hello" {
t.Fatalf("expected 'hello' for single input, got %q", v)
}
// Multiple non-node inputs: type assertion fails, returns zero value
if v := s.Apply(context.Background(), "a", "b"); v != "" {
t.Fatalf("expected zero value for multiple non-node inputs, got %q", v)
}
}
// --- parallelTransporter unit tests ---
func TestParallelTransporter_Dial_NoNodes(t *testing.T) {
tr := &parallelTransporter{nodes: nil}
_, err := tr.Dial(context.Background(), "")
if err == nil {
t.Fatal("expected error with no nodes")
}
}
func TestParallelTransporter_Copy(t *testing.T) {
nodes := []*chain.Node{makeNode(&mockTransport{}), makeNode(&mockTransport{})}
tr := &parallelTransporter{nodes: nodes}
cp := tr.Copy()
cpTr := cp.(*parallelTransporter)
if len(cpTr.nodes) != len(tr.nodes) {
t.Fatalf("copy should have same number of nodes: %d vs %d", len(cpTr.nodes), len(tr.nodes))
}
// Should be a shallow copy of the slice (different backing array)
cpTr.nodes[0] = nil
if tr.nodes[0] == nil {
t.Fatal("copy should not share backing array")
}
}
func TestParallelTransporter_Multiplex(t *testing.T) {
tr := &parallelTransporter{}
if tr.Multiplex() {
t.Fatal("parallel transporter should not multiplex")
}
}
func TestParallelTransporter_Options(t *testing.T) {
tr := &parallelTransporter{}
if opts := tr.Options(); opts == nil {
t.Fatal("expected non-nil options")
}
}
func TestParallelTransporter_Handshake_NonParallelConn(t *testing.T) {
tr := &parallelTransporter{}
c1, c2 := net.Pipe()
defer c1.Close()
defer c2.Close()
conn, err := tr.Handshake(context.Background(), c1)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
// Non-parallelConn should be returned as-is
if conn != c1 {
t.Fatal("expected original conn for non-parallelConn")
}
}
func TestParallelTransporter_Connect_NonParallelConn(t *testing.T) {
tr := &parallelTransporter{}
c1, _ := net.Pipe()
defer c1.Close()
_, err := tr.Connect(context.Background(), c1, "tcp", "addr")
if err == nil {
t.Fatal("expected error for non-parallelConn")
}
}
func TestParallelTransporter_Bind_NonParallelConn(t *testing.T) {
tr := &parallelTransporter{}
c1, _ := net.Pipe()
defer c1.Close()
_, err := tr.Bind(context.Background(), c1, "tcp", "addr")
if err == nil {
t.Fatal("expected error for non-parallelConn")
}
}
// --- parallelConn tests ---
func TestParallelConn_Handshake_Delegates(t *testing.T) {
tr := &parallelTransporter{
nodes: []*chain.Node{makeNode(&mockTransport{})},
}
c1, c2 := net.Pipe()
defer c2.Close()
pc := &parallelConn{Conn: c1, node: makeNode(&mockTransport{})}
conn, err := tr.Handshake(context.Background(), pc)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if conn == nil {
t.Fatal("expected non-nil conn after handshake")
}
conn.Close()
}
func TestParallelConn_Connect_Delegates(t *testing.T) {
tr := &parallelTransporter{
nodes: []*chain.Node{makeNode(&mockTransport{})},
}
c1, c2 := net.Pipe()
defer c2.Close()
pc := &parallelConn{Conn: c1, node: makeNode(&mockTransport{})}
conn, err := tr.Connect(context.Background(), pc, "tcp", "127.0.0.1:80")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if conn == nil {
t.Fatal("expected non-nil conn after connect")
}
conn.Close()
}
func TestParallelConn_Bind_Delegates(t *testing.T) {
tr := &parallelTransporter{
nodes: []*chain.Node{makeNode(&mockTransport{})},
}
c1, c2 := net.Pipe()
defer c2.Close()
pc := &parallelConn{Conn: c1, node: makeNode(&mockTransport{})}
ln, err := tr.Bind(context.Background(), pc, "tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
ln.Close()
}
+53
View File
@@ -0,0 +1,53 @@
package selector
import (
"context"
"testing"
"github.com/go-gost/core/selector"
)
func TestNewSelector(t *testing.T) {
s := NewSelector[string](RoundRobinStrategy[string]())
if s == nil {
t.Fatal("expected non-nil selector")
}
}
func TestDefaultSelector_Select_Empty(t *testing.T) {
s := NewSelector[int](RoundRobinStrategy[int]())
if v := s.Select(context.Background()); v != 0 {
t.Fatalf("expected zero value, got %d", v)
}
}
func TestDefaultSelector_Select_Single(t *testing.T) {
s := NewSelector[int](RoundRobinStrategy[int]())
if v := s.Select(context.Background(), 42); v != 42 {
t.Fatalf("expected 42, got %d", v)
}
}
func TestDefaultSelector_Select_WithFilters(t *testing.T) {
strategy := FIFOStrategy[int]()
filter := selector.Filter[int](BackupFilter[int]())
s := NewSelector[int](strategy, filter)
// Without backup flag, all items pass through; FIFO picks first
v := s.Select(context.Background(), 1, 2, 3)
if v != 1 {
t.Fatalf("expected 1, got %d", v)
}
}
func TestDefaultSelector_FiltersAppliedBeforeStrategy(t *testing.T) {
// Verify that filters narrow the list before strategy picks
strategy := FIFOStrategy[string]()
s := NewSelector[string](strategy, FailFilter[string](1, 0))
// All items are non-markable, so failFilter passes them all through
v := s.Select(context.Background(), "a", "b", "c")
if v != "a" {
t.Fatalf("expected 'a', got %q", v)
}
}
+14
View File
@@ -0,0 +1,14 @@
package selector
import (
"os"
"testing"
corelogger "github.com/go-gost/core/logger"
xlogger "github.com/go-gost/x/logger"
)
func TestMain(m *testing.M) {
corelogger.SetDefault(xlogger.Nop())
os.Exit(m.Run())
}
+5 -4
View File
@@ -93,6 +93,9 @@ type hashStrategy[T any] struct {
mu sync.Mutex mu sync.Mutex
} }
// HashStrategy is a strategy for node selector.
// The node will be selected by hashing the client ID or hash source
// from the context. Falls back to random selection when neither is available.
func HashStrategy[T any]() selector.Strategy[T] { func HashStrategy[T any]() selector.Strategy[T] {
return &hashStrategy[T]{ return &hashStrategy[T]{
r: rand.New(rand.NewSource(time.Now().UnixNano())), r: rand.New(rand.NewSource(time.Now().UnixNano())),
@@ -104,11 +107,9 @@ func (s *hashStrategy[T]) Apply(ctx context.Context, vs ...T) (v T) {
return return
} }
// 打印 基于用户 id 进行 hash 选择最终出口
if clientID := xctx.ClientIDFromContext(ctx); clientID != "" { if clientID := xctx.ClientIDFromContext(ctx); clientID != "" {
idStr := string(clientID) value := uint64(crc32.ChecksumIEEE([]byte(string(clientID))))
value := uint64(crc32.ChecksumIEEE([]byte(idStr))) logger.Default().Tracef("hash client %s %d", clientID, value)
logger.Default().Infof("【命中】从 xctx 拿到 ID: %s, Hash: %d", idStr, value)
return vs[value%uint64(len(vs))] return vs[value%uint64(len(vs))]
} }
+271
View File
@@ -0,0 +1,271 @@
package selector
import (
"context"
"hash/crc32"
"sync"
"testing"
"github.com/go-gost/core/metadata"
"github.com/go-gost/core/selector"
xctx "github.com/go-gost/x/ctx"
xmd "github.com/go-gost/x/metadata"
)
// --- RoundRobinStrategy ---
func TestRoundRobinStrategy_Empty(t *testing.T) {
s := RoundRobinStrategy[int]()
if v := s.Apply(context.Background()); v != 0 {
t.Fatalf("expected zero value, got %d", v)
}
}
func TestRoundRobinStrategy_Single(t *testing.T) {
s := RoundRobinStrategy[int]()
for i := 0; i < 5; i++ {
if v := s.Apply(context.Background(), 42); v != 42 {
t.Fatalf("iteration %d: expected 42, got %d", i, v)
}
}
}
func TestRoundRobinStrategy_Sequence(t *testing.T) {
s := RoundRobinStrategy[int]()
items := []int{10, 20, 30}
// Should cycle through 10, 20, 30, 10, 20, 30, ...
for i := 0; i < 9; i++ {
v := s.Apply(context.Background(), items...)
expected := items[i%3]
if v != expected {
t.Fatalf("iteration %d: expected %d, got %d", i, expected, v)
}
}
}
func TestRoundRobinStrategy_Concurrent(t *testing.T) {
s := RoundRobinStrategy[int]()
items := []int{0, 1, 2, 3, 4}
const goroutines = 100
var wg sync.WaitGroup
results := make(chan int, goroutines)
for i := 0; i < goroutines; i++ {
wg.Add(1)
go func() {
defer wg.Done()
results <- s.Apply(context.Background(), items...)
}()
}
wg.Wait()
close(results)
// Every result should be a valid item
for v := range results {
if v < 0 || v > 4 {
t.Fatalf("unexpected value %d", v)
}
}
}
// --- FIFOStrategy ---
func TestFIFOStrategy_Empty(t *testing.T) {
s := FIFOStrategy[string]()
if v := s.Apply(context.Background()); v != "" {
t.Fatalf("expected zero value, got %q", v)
}
}
func TestFIFOStrategy_AlwaysFirst(t *testing.T) {
s := FIFOStrategy[int]()
for i := 0; i < 10; i++ {
v := s.Apply(context.Background(), 100, 200, 300)
if v != 100 {
t.Fatalf("iteration %d: expected 100, got %d", i, v)
}
}
}
// --- RandomStrategy ---
func TestRandomStrategy_Empty(t *testing.T) {
s := RandomStrategy[int]()
if v := s.Apply(context.Background()); v != 0 {
t.Fatalf("expected zero value, got %d", v)
}
}
func TestRandomStrategy_Single(t *testing.T) {
s := RandomStrategy[int]()
for i := 0; i < 10; i++ {
if v := s.Apply(context.Background(), 7); v != 7 {
t.Fatalf("expected 7, got %d", v)
}
}
}
func TestRandomStrategy_Distribution(t *testing.T) {
s := RandomStrategy[int]()
items := []int{0, 1, 2}
counts := make(map[int]int)
const n = 3000
for i := 0; i < n; i++ {
v := s.Apply(context.Background(), items...)
counts[v]++
}
// With uniform weights, each should get roughly n/3
for _, item := range items {
if counts[item] < n/len(items)/2 {
t.Fatalf("item %d underrepresented: %d/%d", item, counts[item], n)
}
}
}
// weightedItem implements metadata.Metadatable for weight testing.
type weightedItem struct {
md metadata.Metadata
}
func (w *weightedItem) Metadata() metadata.Metadata { return w.md }
func TestRandomStrategy_Weighted(t *testing.T) {
s := RandomStrategy[*weightedItem]()
heavy := &weightedItem{md: xmd.NewMetadata(map[string]any{"weight": 100})}
light := &weightedItem{md: xmd.NewMetadata(map[string]any{"weight": 1})}
heavyCount := 0
const n = 2000
for i := 0; i < n; i++ {
v := s.Apply(context.Background(), heavy, light)
if v == heavy {
heavyCount++
}
}
// Heavy should win ~99% of the time with weight 100 vs 1
if heavyCount < int(float64(n)*0.9) {
t.Fatalf("heavy item underrepresented: %d/%d", heavyCount, n)
}
}
func TestRandomStrategy_ZeroWeight(t *testing.T) {
s := RandomStrategy[int]()
// Zero/negative weights default to 1, so all items should be selectable
v := s.Apply(context.Background(), 1, 2, 3)
if v < 1 || v > 3 {
t.Fatalf("expected 1-3, got %d", v)
}
}
// --- HashStrategy ---
func TestHashStrategy_Empty(t *testing.T) {
s := HashStrategy[int]()
if v := s.Apply(context.Background()); v != 0 {
t.Fatalf("expected zero value, got %d", v)
}
}
func TestHashStrategy_ClientID(t *testing.T) {
s := HashStrategy[int]()
items := []int{10, 20, 30, 40, 50}
ctx := xctx.ContextWithClientID(context.Background(), "test-client")
// Same client ID should always map to the same item
v1 := s.Apply(ctx, items...)
v2 := s.Apply(ctx, items...)
if v1 != v2 {
t.Fatalf("hash strategy should be deterministic for same client ID: %d != %d", v1, v2)
}
// Verify it matches manual CRC32 calculation
expectedIdx := uint64(crc32.ChecksumIEEE([]byte("test-client"))) % uint64(len(items))
if v1 != items[expectedIdx] {
t.Fatalf("expected items[%d]=%d, got %d", expectedIdx, items[expectedIdx], v1)
}
}
func TestHashStrategy_HashSource(t *testing.T) {
s := HashStrategy[int]()
items := []int{10, 20, 30}
ctx := xctx.ContextWithHash(context.Background(), &xctx.Hash{Source: "my-hash-key"})
v1 := s.Apply(ctx, items...)
v2 := s.Apply(ctx, items...)
if v1 != v2 {
t.Fatalf("hash strategy should be deterministic for same hash source: %d != %d", v1, v2)
}
}
func TestHashStrategy_ClientIDPriorityOverHash(t *testing.T) {
s := HashStrategy[int]()
items := []int{10, 20, 30, 40}
ctxClient := xctx.ContextWithClientID(context.Background(), "clientA")
ctxBoth := xctx.ContextWithHash(
xctx.ContextWithClientID(context.Background(), "clientA"),
&xctx.Hash{Source: "different-hash"},
)
vClient := s.Apply(ctxClient, items...)
vBoth := s.Apply(ctxBoth, items...)
// Client ID should take priority over hash source
if vClient != vBoth {
t.Fatalf("client ID should take priority: %d != %d", vClient, vBoth)
}
}
func TestHashStrategy_FallbackRandom(t *testing.T) {
s := HashStrategy[int]()
items := []int{0, 1, 2}
counts := make(map[int]int)
const n = 3000
for i := 0; i < n; i++ {
v := s.Apply(context.Background(), items...)
counts[v]++
}
// Without client ID or hash, falls back to random
for _, item := range items {
if counts[item] < n/len(items)/2 {
t.Fatalf("item %d underrepresented in random fallback: %d/%d", item, counts[item], n)
}
}
}
func TestHashStrategy_DifferentClientIDs(t *testing.T) {
s := HashStrategy[int]()
items := make([]int, 100)
for i := range items {
items[i] = i
}
selected := make(map[int]bool)
for i := 0; i < 100; i++ {
ctx := xctx.ContextWithClientID(context.Background(), xctx.ClientID(string(rune('A'+i))))
v := s.Apply(ctx, items...)
selected[v] = true
}
// Different client IDs should spread across multiple items
if len(selected) < 10 {
t.Fatalf("expected distribution across many items, got %d unique selections", len(selected))
}
}
// --- Interface compliance ---
func TestStrategyInterfaceCompliance(t *testing.T) {
// Verify all strategies implement selector.Strategy
_ = selector.Strategy[int](RoundRobinStrategy[int]())
_ = selector.Strategy[int](RandomStrategy[int]())
_ = selector.Strategy[int](FIFOStrategy[int]())
_ = selector.Strategy[int](HashStrategy[int]())
_ = selector.Strategy[int](ParallelStrategy[int]())
}
+7
View File
@@ -10,24 +10,30 @@ type randomWeightedItem[T any] struct {
weight int weight int
} }
// RandomWeighted is a weighted random selector.
// Items are selected randomly with probability proportional to their weight.
type RandomWeighted[T any] struct { type RandomWeighted[T any] struct {
items []*randomWeightedItem[T] items []*randomWeightedItem[T]
sum int sum int
r *rand.Rand r *rand.Rand
} }
// NewRandomWeighted creates a new RandomWeighted selector.
func NewRandomWeighted[T any]() *RandomWeighted[T] { func NewRandomWeighted[T any]() *RandomWeighted[T] {
return &RandomWeighted[T]{ return &RandomWeighted[T]{
r: rand.New(rand.NewSource(time.Now().UnixNano())), r: rand.New(rand.NewSource(time.Now().UnixNano())),
} }
} }
// Add adds an item with the given weight to the selector.
func (rw *RandomWeighted[T]) Add(item T, weight int) { func (rw *RandomWeighted[T]) Add(item T, weight int) {
ri := &randomWeightedItem[T]{item: item, weight: weight} ri := &randomWeightedItem[T]{item: item, weight: weight}
rw.items = append(rw.items, ri) rw.items = append(rw.items, ri)
rw.sum += weight rw.sum += weight
} }
// Next selects an item randomly based on weights.
// Returns the zero value if no items are added.
func (rw *RandomWeighted[T]) Next() (v T) { func (rw *RandomWeighted[T]) Next() (v T) {
if len(rw.items) == 0 { if len(rw.items) == 0 {
return return
@@ -46,6 +52,7 @@ func (rw *RandomWeighted[T]) Next() (v T) {
return rw.items[len(rw.items)-1].item return rw.items[len(rw.items)-1].item
} }
// Reset clears all items and weights.
func (rw *RandomWeighted[T]) Reset() { func (rw *RandomWeighted[T]) Reset() {
rw.items = nil rw.items = nil
rw.sum = 0 rw.sum = 0
+130
View File
@@ -0,0 +1,130 @@
package selector
import (
"testing"
)
func TestRandomWeighted_Empty(t *testing.T) {
rw := NewRandomWeighted[int]()
if v := rw.Next(); v != 0 {
t.Fatalf("expected zero value, got %d", v)
}
}
func TestRandomWeighted_Single(t *testing.T) {
rw := NewRandomWeighted[int]()
rw.Add(42, 10)
for i := 0; i < 100; i++ {
if v := rw.Next(); v != 42 {
t.Fatalf("expected 42, got %d", v)
}
}
}
func TestRandomWeighted_EqualWeights(t *testing.T) {
rw := NewRandomWeighted[int]()
rw.Add(1, 1)
rw.Add(2, 1)
rw.Add(3, 1)
counts := make(map[int]int)
const n = 3000
for i := 0; i < n; i++ {
counts[rw.Next()]++
}
for _, item := range []int{1, 2, 3} {
if counts[item] < n/3/2 {
t.Fatalf("item %d underrepresented: %d/%d", item, counts[item], n)
}
}
}
func TestRandomWeighted_WeightedDistribution(t *testing.T) {
rw := NewRandomWeighted[string]()
rw.Add("heavy", 99)
rw.Add("light", 1)
heavyCount := 0
const n = 5000
for i := 0; i < n; i++ {
if rw.Next() == "heavy" {
heavyCount++
}
}
// ~99% should be heavy
if heavyCount < int(float64(n)*0.95) {
t.Fatalf("heavy underrepresented: %d/%d", heavyCount, n)
}
}
func TestRandomWeighted_ZeroWeight(t *testing.T) {
rw := NewRandomWeighted[int]()
rw.Add(1, 0)
rw.Add(2, 0)
// sum=0, should return zero value
if v := rw.Next(); v != 0 {
t.Fatalf("expected zero value with sum=0, got %d", v)
}
}
func TestRandomWeighted_Reset(t *testing.T) {
rw := NewRandomWeighted[int]()
rw.Add(1, 10)
rw.Add(2, 20)
rw.Reset()
if v := rw.Next(); v != 0 {
t.Fatalf("expected zero value after reset, got %d", v)
}
// Can add again after reset
rw.Add(3, 1)
if v := rw.Next(); v != 3 {
t.Fatalf("expected 3 after re-add, got %d", v)
}
}
func TestRandomWeighted_MultipleWeights(t *testing.T) {
rw := NewRandomWeighted[int]()
rw.Add(0, 1)
rw.Add(1, 2)
rw.Add(2, 3)
// Total weight = 6
counts := make(map[int]int)
const n = 6000
for i := 0; i < n; i++ {
counts[rw.Next()]++
}
// Item 2 should appear roughly 2x more than item 0
ratio := float64(counts[2]) / float64(counts[0])
if ratio < 1.5 || ratio > 3.5 {
t.Fatalf("expected ratio ~2.0 for item2/item0, got %.2f (counts: %v)", ratio, counts)
}
}
func TestRandomWeighted_StringItems(t *testing.T) {
rw := NewRandomWeighted[string]()
rw.Add("a", 1)
rw.Add("b", 1)
v := rw.Next()
if v != "a" && v != "b" {
t.Fatalf("expected 'a' or 'b', got %q", v)
}
}
func TestRandomWeighted_StructItems(t *testing.T) {
type item struct{ Name string }
rw := NewRandomWeighted[item]()
rw.Add(item{Name: "x"}, 1)
rw.Add(item{Name: "y"}, 1)
v := rw.Next()
if v.Name != "x" && v.Name != "y" {
t.Fatalf("unexpected item: %+v", v)
}
}