p2p/enode: add tests for source naming

This commit is contained in:
Felix Lange 2025-05-07 21:13:10 +02:00
parent 3e8a5f318a
commit cb3c69b732
2 changed files with 50 additions and 2 deletions

View file

@ -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 {

View file

@ -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 {