go-libp2p-pubsub/randomsub_test.go

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...)
	if err != nil {
		panic(err)
	}
	return ps
}

func getRandomsubs(ctx context.Context, hs []host.Host, size int, opts ...Option) []*PubSub {
	var psubs []*PubSub
	for _, h := range hs {
		psubs = append(psubs, getRandomsub(ctx, h, size, opts...))
	}
	return psubs
}

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
	}
}

func TestRandomsubSmall(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	hosts := getNetHosts(t, ctx, 10)
	psubs := getRandomsubs(ctx, hosts, 10)

	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)

	count := 0
	for i := 0; i < 10; i++ {
		msg := []byte(fmt.Sprintf("message %d", i))
		psubs[i].Publish("test", msg)

		for _, sub := range subs {
			if tryReceive(sub) != nil {
				count++
			}
		}
	}

	if count < 9*len(hosts) {
		t.Fatalf("received too few messages; expected at least %d but got %d", 9*len(hosts), count)
	}
}

func TestRandomsubBig(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	hosts := getNetHosts(t, ctx, 50)
	psubs := getRandomsubs(ctx, hosts, 50)

	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)

	count := 0
	for i := 0; i < 10; i++ {
		msg := []byte(fmt.Sprintf("message %d", i))
		psubs[i].Publish("test", msg)

		for _, sub := range subs {
			if tryReceive(sub) != nil {
				count++
			}
		}
	}

	if count < 9*len(hosts) {
		t.Fatalf("received too few messages; expected at least %d but got %d", 9*len(hosts), count)
	}
}

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)
	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)

	count := 0
	for i := 0; i < 10; i++ {
		msg := []byte(fmt.Sprintf("message %d", i))
		psubs[i].Publish("test", msg)

		for _, sub := range subs {
			if tryReceive(sub) != nil {
				count++
			}
		}
	}

	if count < 9*len(hosts) {
		t.Fatalf("received too few messages; expected at least %d but got %d", 9*len(hosts), count)
	}
}

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)
	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")
	}
}