go-libp2p-pubsub/randomsub_test.go

195 lines
3.8 KiB
Go
Raw Normal View History

2020-04-23 16:37:12 +00:00
package pubsub
import (
"context"
"fmt"
"testing"
"time"
"github.com/libp2p/go-libp2p-core/host"
)
func getRandomsub(ctx context.Context, h host.Host, size int, opts ...Option) *PubSub {
ps, err := NewRandomSub(ctx, h, size, opts...)
2020-04-23 16:37:12 +00:00
if err != nil {
panic(err)
}
return ps
}
func getRandomsubs(ctx context.Context, hs []host.Host, size int, opts ...Option) []*PubSub {
2020-04-23 16:37:12 +00:00
var psubs []*PubSub
for _, h := range hs {
psubs = append(psubs, getRandomsub(ctx, h, size, opts...))
2020-04-23 16:37:12 +00:00
}
return psubs
}
2020-04-23 17:35:14 +00:00
func tryReceive(sub *Subscription) *Message {
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
m, err := sub.Next(ctx)
if err != nil {
return nil
} else {
return m
}
}
2020-04-23 16:37:12 +00:00
func TestRandomsubSmall(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
hosts := getNetHosts(t, ctx, 10)
psubs := getRandomsubs(ctx, hosts, 10)
2020-04-23 16:37:12 +00:00
connectAll(t, hosts)
var subs []*Subscription
for _, ps := range psubs {
sub, err := ps.Subscribe("test")
if err != nil {
t.Fatal(err)
}
subs = append(subs, sub)
}
time.Sleep(time.Second)
2020-04-23 17:35:14 +00:00
count := 0
2020-04-23 16:37:12 +00:00
for i := 0; i < 10; i++ {
msg := []byte(fmt.Sprintf("message %d", i))
psubs[i].Publish("test", msg)
for _, sub := range subs {
2020-04-23 17:35:14 +00:00
if tryReceive(sub) != nil {
count++
}
2020-04-23 16:37:12 +00:00
}
}
2020-04-23 17:35:14 +00:00
if count < 7*len(hosts) {
2020-04-23 17:35:14 +00:00
t.Fatalf("received too few messages; expected at least %d but got %d", 9*len(hosts), count)
}
2020-04-23 16:37:12 +00:00
}
func TestRandomsubBig(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
hosts := getNetHosts(t, ctx, 50)
psubs := getRandomsubs(ctx, hosts, 50)
2020-04-23 16:37:12 +00:00
connectSome(t, hosts, 12)
var subs []*Subscription
for _, ps := range psubs {
sub, err := ps.Subscribe("test")
if err != nil {
t.Fatal(err)
}
subs = append(subs, sub)
}
time.Sleep(time.Second)
2020-04-23 17:35:14 +00:00
count := 0
2020-04-23 16:37:12 +00:00
for i := 0; i < 10; i++ {
msg := []byte(fmt.Sprintf("message %d", i))
psubs[i].Publish("test", msg)
for _, sub := range subs {
2020-04-23 17:35:14 +00:00
if tryReceive(sub) != nil {
count++
}
2020-04-23 16:37:12 +00:00
}
}
2020-04-23 17:35:14 +00:00
if count < 7*len(hosts) {
2020-04-23 17:35:14 +00:00
t.Fatalf("received too few messages; expected at least %d but got %d", 9*len(hosts), count)
}
2020-04-23 16:37:12 +00:00
}
func TestRandomsubMixed(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
hosts := getNetHosts(t, ctx, 40)
fsubs := getPubsubs(ctx, hosts[:10])
rsubs := getRandomsubs(ctx, hosts[10:], 30)
2020-04-23 16:37:12 +00:00
psubs := append(fsubs, rsubs...)
connectSome(t, hosts, 12)
var subs []*Subscription
for _, ps := range psubs {
sub, err := ps.Subscribe("test")
if err != nil {
t.Fatal(err)
}
subs = append(subs, sub)
}
time.Sleep(time.Second)
2020-04-23 17:35:14 +00:00
count := 0
2020-04-23 16:37:12 +00:00
for i := 0; i < 10; i++ {
msg := []byte(fmt.Sprintf("message %d", i))
psubs[i].Publish("test", msg)
for _, sub := range subs {
2020-04-23 17:35:14 +00:00
if tryReceive(sub) != nil {
count++
}
2020-04-23 16:37:12 +00:00
}
}
2020-04-23 17:35:14 +00:00
if count < 7*len(hosts) {
2020-04-23 17:35:14 +00:00
t.Fatalf("received too few messages; expected at least %d but got %d", 9*len(hosts), count)
}
2020-04-23 16:37:12 +00:00
}
func TestRandomsubEnoughPeers(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
hosts := getNetHosts(t, ctx, 40)
fsubs := getPubsubs(ctx, hosts[:10])
rsubs := getRandomsubs(ctx, hosts[10:], 30)
2020-04-23 16:37:12 +00:00
psubs := append(fsubs, rsubs...)
connectSome(t, hosts, 12)
var subs []*Subscription
for _, ps := range psubs {
sub, err := ps.Subscribe("test")
if err != nil {
t.Fatal(err)
}
subs = append(subs, sub)
}
time.Sleep(time.Second)
res := make(chan bool, 1)
rsubs[0].eval <- func() {
rs := rsubs[0].rt.(*RandomSubRouter)
res <- rs.EnoughPeers("test", 0)
}
enough := <-res
if !enough {
t.Fatal("expected enough peers")
}
rsubs[0].eval <- func() {
rs := rsubs[0].rt.(*RandomSubRouter)
res <- rs.EnoughPeers("test", 100)
}
enough = <-res
if !enough {
t.Fatal("expected enough peers")
}
}