diff --git a/p2p/node/pubsubManager/gossipsub_test.go b/p2p/node/pubsubManager/gossipsub_test.go index 016835358..012416666 100644 --- a/p2p/node/pubsubManager/gossipsub_test.go +++ b/p2p/node/pubsubManager/gossipsub_test.go @@ -4,6 +4,7 @@ import ( "context" "sync" "testing" + "time" "go.uber.org/mock/gomock" @@ -139,9 +140,11 @@ func TestPubsubManager(t *testing.T) { func TestMultipleRequests(t *testing.T) { // Number of requests to test - n := 10 + n := 100 - ctx := context.Background() + // Use a context with timeout to avoid hanging indefinitely + ctx, cancel := context.WithTimeout(context.Background(), 1*time.Minute) + defer cancel() mockHost, mockPeerStore, privKey, peerID := setup(t) @@ -180,10 +183,14 @@ func TestMultipleRequests(t *testing.T) { } } - //BROADCAST - testCh := make(chan interface{}, n) + // BROADCAST + testCh := make(chan interface{}, n*len(topics)) ps.SetReceiveHandler(func(receivedFrom peer.ID, msgId string, msgTopic string, data interface{}, location common.Location) { - testCh <- data + select { + case testCh <- data: + case <-ctx.Done(): + t.Error("context done before send messageId: ", msgId) + } }) var messages []interface{} @@ -215,23 +222,29 @@ func TestMultipleRequests(t *testing.T) { } // VERIFY - receivedMessages := make([]interface{}, 0, n) + var mu sync.Mutex + receivedMessages := make([]interface{}, 0, n*len(topics)) for i := 0; i < (n * len(topics)); i++ { wg.Add(1) - go func() { + go func(j int) { defer wg.Done() - receivedMessage := <-testCh - receivedMessages = append(receivedMessages, receivedMessage) - }() + select { + case receivedMessage := <-testCh: + mu.Lock() + receivedMessages = append(receivedMessages, receivedMessage) + mu.Unlock() + case <-ctx.Done(): + t.Error("context done before receive message at index: ", j) + } + }(i) } wg.Wait() // Ensure all broadcasted messages were received - require.Len(t, receivedMessages, len(messages), "The number of received messages does not match the number of broadcasted messages") + require.Len(t, receivedMessages, len(messages), "The number of received messages does not match the number of broadcasted messages. expected: %d, got: %d", len(messages), len(receivedMessages)) - // UNSUBSCRIBE All - ps.UnsubscribeAll() + ps.Stop() if len(ps.GetTopics()) != 0 { t.Fatal("Topic should be empty after unsubscribe") }