github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/pubsub/libp2p/pubsub_test.go (about)

     1  package libp2p
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	libp2p_host "github.com/filecoin-project/bacalhau/pkg/libp2p"
     9  	"github.com/filecoin-project/bacalhau/pkg/pubsub"
    10  	libp2p_pubsub "github.com/libp2p/go-libp2p-pubsub"
    11  	"github.com/libp2p/go-libp2p/core/host"
    12  	"github.com/rs/zerolog/log"
    13  	"github.com/stretchr/testify/suite"
    14  )
    15  
    16  const testTopic = "test-topic"
    17  
    18  type PubSubSuite struct {
    19  	suite.Suite
    20  	node1       *PubSub[string]
    21  	node2       *PubSub[string]
    22  	subscriber1 *pubsub.InMemorySubscriber[string]
    23  	subscriber2 *pubsub.InMemorySubscriber[string]
    24  }
    25  
    26  func (s *PubSubSuite) SetupSuite() {
    27  	n1, h1 := s.createPubSub(false)
    28  	s.node1 = n1
    29  	s.node2, _ = s.createPubSub(true, h1)
    30  
    31  	s.subscriber1 = pubsub.NewInMemorySubscriber[string]()
    32  	s.subscriber2 = pubsub.NewInMemorySubscriber[string]()
    33  	s.NoError(s.node1.Subscribe(context.Background(), s.subscriber1))
    34  	s.NoError(s.node2.Subscribe(context.Background(), s.subscriber2))
    35  
    36  	// wait for up to 10 seconds (5 loops with 2 seconds each) for nodes to discover each other
    37  	var s1, s2 bool
    38  	for i := 0; i < 5; i++ {
    39  		s.NoError(s.node1.Publish(context.Background(), "ping"))
    40  		s1, s2 = s.waitForMessage("ping", 2*time.Second, true, true)
    41  		if s1 || s2 {
    42  			// still one of the subscribers is waiting for the message
    43  			continue
    44  		}
    45  	}
    46  	if s1 {
    47  		s.FailNow("subscriber 1 didn't receive initialization message")
    48  	}
    49  	if s2 {
    50  		s.FailNow("subscriber 2 didn't receive initialization message")
    51  	}
    52  	log.Debug().Msg("libp2p pubsub suite is ready")
    53  }
    54  
    55  func (s *PubSubSuite) TearDownSuite() {
    56  	s.NoError(s.node1.Close(context.Background()))
    57  	s.NoError(s.node2.Close(context.Background()))
    58  }
    59  
    60  func (s *PubSubSuite) createPubSub(ignoreLocal bool, peers ...host.Host) (*PubSub[string], host.Host) {
    61  	h, err := libp2p_host.NewHostForTest(context.Background(), peers...)
    62  	s.NoError(err)
    63  
    64  	gossipSub, err := libp2p_pubsub.NewGossipSub(context.Background(), h)
    65  	s.NoError(err)
    66  
    67  	pubSub, err := NewPubSub[string](PubSubParams{
    68  		Host:        h,
    69  		TopicName:   testTopic,
    70  		PubSub:      gossipSub,
    71  		IgnoreLocal: ignoreLocal,
    72  	})
    73  	s.NoError(err)
    74  
    75  	return pubSub, h
    76  }
    77  
    78  func TestPubSubSuite(t *testing.T) {
    79  	suite.Run(t, new(PubSubSuite))
    80  }
    81  
    82  func (s *PubSubSuite) TestPubSub() {
    83  	msg := "TestPubSub"
    84  	s.NoError(s.node1.Publish(context.Background(), msg))
    85  	s.waitForMessage(msg, 10*time.Second, true, true)
    86  }
    87  
    88  func (s *PubSubSuite) TestPubSub_IgnoreLocal() {
    89  	// node2 is ignoring local messages, so it should not receive the message
    90  	msg := "TestPubSub_IgnoreLocal"
    91  	s.NoError(s.node2.Publish(context.Background(), msg))
    92  	s.waitForMessage(msg, 10*time.Second, true, false)
    93  	s.Empty(s.subscriber2.Events())
    94  }
    95  
    96  func (s *PubSubSuite) waitForMessage(msg string, duration time.Duration, checkSubscriber1, checkSubscriber2 bool) (bool, bool) {
    97  	waitUntil := time.Now().Add(duration)
    98  	checkSubscriber := func(subscriber *pubsub.InMemorySubscriber[string]) bool {
    99  		events := subscriber.Events()
   100  		if len(events) == 0 {
   101  			return false
   102  		}
   103  		s.Equal([]string{msg}, events)
   104  		return true
   105  	}
   106  
   107  	for time.Now().Before(waitUntil) && (checkSubscriber1 || checkSubscriber2) {
   108  		time.Sleep(100 * time.Millisecond)
   109  		if checkSubscriber1 && checkSubscriber(s.subscriber1) {
   110  			checkSubscriber1 = false
   111  		}
   112  		if checkSubscriber2 && checkSubscriber(s.subscriber2) {
   113  			checkSubscriber2 = false
   114  		}
   115  	}
   116  
   117  	return checkSubscriber1, checkSubscriber1
   118  }