p2p/enode: add tests for source naming

p2p/enode/iter.go

@@ -38,7 +38,7 @@ type SourceIterator interface {
 }
 
 // WithSource attaches a 'source name' to an iterator.
-func WithSource(name string, it Iterator) SourceIterator {
+func WithSourceName(name string, it Iterator) SourceIterator {
 	return sourceIter{it, name}
 }
 
@@ -46,7 +46,7 @@ func ensureSourceIter(it Iterator) SourceIterator {
 	if si, ok := it.(SourceIterator); ok {
 		return si
 	}
-	return WithSource("", it)
+	return WithSourceName("", it)
 }
 
 type sourceIter struct {

p2p/enode/iter_test.go

@@ -19,6 +19,7 @@ package enode
 import (
 	"encoding/binary"
 	"runtime"
+	"slices"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -183,6 +184,53 @@ func TestFairMixRemoveSource(t *testing.T) {
 	}
 }
 
+// This checks that FairMix correctly returns the name of the source that produced the node.
+func TestFairMixSourceName(t *testing.T) {
+	nodes := make([]*Node, 6)
+	for i := range nodes {
+		nodes[i] = testNode(uint64(i), uint64(i))
+	}
+	mix := NewFairMix(-1)
+	mix.AddSource(WithSourceName("s1", IterNodes(nodes[0:2])))
+	mix.AddSource(WithSourceName("s2", IterNodes(nodes[2:4])))
+	mix.AddSource(WithSourceName("s3", IterNodes(nodes[4:6])))
+
+	var names []string
+	for range nodes {
+		mix.Next()
+		names = append(names, mix.NodeSource())
+	}
+	want := []string{"s2", "s3", "s1", "s2", "s3", "s1"}
+	if !slices.Equal(names, want) {
+		t.Fatalf("wrong names: %v", names)
+	}
+}
+
+// This checks that FairMix returns the name of the source that produced the node,
+// even when FairMix instances are nested.
+func TestFairMixNestedSourceName(t *testing.T) {
+	nodes := make([]*Node, 6)
+	for i := range nodes {
+		nodes[i] = testNode(uint64(i), uint64(i))
+	}
+	mix := NewFairMix(-1)
+	mix.AddSource(WithSourceName("s1", IterNodes(nodes[0:2])))
+	submix := NewFairMix(-1)
+	submix.AddSource(WithSourceName("s2", IterNodes(nodes[2:4])))
+	submix.AddSource(WithSourceName("s3", IterNodes(nodes[4:6])))
+	mix.AddSource(submix)
+
+	var names []string
+	for range nodes {
+		mix.Next()
+		names = append(names, mix.NodeSource())
+	}
+	want := []string{"s3", "s1", "s2", "s1", "s3", "s2"}
+	if !slices.Equal(names, want) {
+		t.Fatalf("wrong names: %v", names)
+	}
+}
+
 type blockingIter chan struct{}
 
 func (it blockingIter) Next() bool {