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:
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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 ¶llelStrategy[T]{}
|
return ¶llelStrategy[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(¶llelTransporter{nodes: nodes}))
|
vn := chain.NewNode("parallel", "", chain.TransportNodeOption(¶llelTransporter{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 ¶llelTransporter{
|
return ¶llelTransporter{
|
||||||
nodes: append([]*chain.Node(nil), tr.nodes...),
|
nodes: append([]*chain.Node(nil), tr.nodes...),
|
||||||
|
|||||||
@@ -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 := ¶llelTransporter{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 := ¶llelTransporter{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 := ¶llelTransporter{}
|
||||||
|
if tr.Multiplex() {
|
||||||
|
t.Fatal("parallel transporter should not multiplex")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParallelTransporter_Options(t *testing.T) {
|
||||||
|
tr := ¶llelTransporter{}
|
||||||
|
if opts := tr.Options(); opts == nil {
|
||||||
|
t.Fatal("expected non-nil options")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestParallelTransporter_Handshake_NonParallelConn(t *testing.T) {
|
||||||
|
tr := ¶llelTransporter{}
|
||||||
|
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 := ¶llelTransporter{}
|
||||||
|
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 := ¶llelTransporter{}
|
||||||
|
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 := ¶llelTransporter{
|
||||||
|
nodes: []*chain.Node{makeNode(&mockTransport{})},
|
||||||
|
}
|
||||||
|
|
||||||
|
c1, c2 := net.Pipe()
|
||||||
|
defer c2.Close()
|
||||||
|
|
||||||
|
pc := ¶llelConn{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 := ¶llelTransporter{
|
||||||
|
nodes: []*chain.Node{makeNode(&mockTransport{})},
|
||||||
|
}
|
||||||
|
|
||||||
|
c1, c2 := net.Pipe()
|
||||||
|
defer c2.Close()
|
||||||
|
|
||||||
|
pc := ¶llelConn{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 := ¶llelTransporter{
|
||||||
|
nodes: []*chain.Node{makeNode(&mockTransport{})},
|
||||||
|
}
|
||||||
|
|
||||||
|
c1, c2 := net.Pipe()
|
||||||
|
defer c2.Close()
|
||||||
|
|
||||||
|
pc := ¶llelConn{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()
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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())
|
||||||
|
}
|
||||||
@@ -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))]
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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]())
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user