go-waku/waku/v2/broadcast_test.go

package v2

import (
	"context"
	"sync"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/waku-org/go-waku/waku/v2/protocol"
	"github.com/waku-org/go-waku/waku/v2/protocol/pb"
	"github.com/waku-org/go-waku/waku/v2/utils"
)

// Adapted from https://github.com/dustin/go-broadcast/commit/f664265f5a662fb4d1df7f3533b1e8d0e0277120
// by Dustin Sallings (c) 2013, which was released under MIT license
func TestBroadcast(t *testing.T) {
	wg := sync.WaitGroup{}

	b := NewBroadcaster(100)
	require.NoError(t, b.Start(context.Background()))
	defer b.Stop()

	for i := 0; i < 5; i++ {
		wg.Add(1)

		cch := make(chan *protocol.Envelope)
		b.Register(nil, cch)
		go func() {
			defer wg.Done()
			defer b.Unregister(nil, cch)
			<-cch
		}()

	}

	env := protocol.NewEnvelope(&pb.WakuMessage{}, utils.GetUnixEpoch(), "abc")
	b.Submit(env)

	wg.Wait()
}

func TestBroadcastWait(t *testing.T) {
	wg := sync.WaitGroup{}

	b := NewBroadcaster(100)
	require.NoError(t, b.Start(context.Background()))
	defer b.Stop()

	for i := 0; i < 5; i++ {
		wg.Add(1)

		cch := make(chan *protocol.Envelope)
		<-b.WaitRegister(nil, cch)

		go func() {
			defer wg.Done()

			<-cch
			<-b.WaitUnregister(nil, cch)
		}()

	}

	env := protocol.NewEnvelope(&pb.WakuMessage{}, utils.GetUnixEpoch(), "abc")
	b.Submit(env)

	wg.Wait()
}

func TestBroadcastCleanup(t *testing.T) {
	b := NewBroadcaster(100)
	require.NoError(t, b.Start(context.Background()))
	topic := "test"
	b.Register(&topic, make(chan *protocol.Envelope))
	b.Stop()
}