github.com/prysmaticlabs/prysm@v1.4.4/beacon-chain/p2p/sender_test.go (about)

     1  package p2p
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
    10  
    11  	"github.com/libp2p/go-libp2p-core/network"
    12  	testp2p "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
    13  	"github.com/prysmaticlabs/prysm/shared/testutil"
    14  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    15  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    16  	"google.golang.org/protobuf/proto"
    17  )
    18  
    19  func TestService_Send(t *testing.T) {
    20  	p1 := testp2p.NewTestP2P(t)
    21  	p2 := testp2p.NewTestP2P(t)
    22  	p1.Connect(p2)
    23  
    24  	svc := &Service{
    25  		host: p1.BHost,
    26  		cfg:  &Config{},
    27  	}
    28  
    29  	msg := &pb.Fork{
    30  		CurrentVersion:  []byte("fooo"),
    31  		PreviousVersion: []byte("barr"),
    32  		Epoch:           55,
    33  	}
    34  
    35  	// Register external listener which will repeat the message back.
    36  	var wg sync.WaitGroup
    37  	wg.Add(1)
    38  	topic := "/testing/1"
    39  	RPCTopicMappings[topic] = new(pb.Fork)
    40  	defer func() {
    41  		delete(RPCTopicMappings, topic)
    42  	}()
    43  	p2.SetStreamHandler(topic+"/ssz_snappy", func(stream network.Stream) {
    44  		rcvd := &pb.Fork{}
    45  		require.NoError(t, svc.Encoding().DecodeWithMaxLength(stream, rcvd))
    46  		_, err := svc.Encoding().EncodeWithMaxLength(stream, rcvd)
    47  		require.NoError(t, err)
    48  		assert.NoError(t, stream.Close())
    49  		wg.Done()
    50  	})
    51  
    52  	stream, err := svc.Send(context.Background(), msg, "/testing/1", p2.BHost.ID())
    53  	require.NoError(t, err)
    54  
    55  	testutil.WaitTimeout(&wg, 1*time.Second)
    56  
    57  	rcvd := &pb.Fork{}
    58  	require.NoError(t, svc.Encoding().DecodeWithMaxLength(stream, rcvd))
    59  	if !proto.Equal(rcvd, msg) {
    60  		t.Errorf("Expected identical message to be received. got %v want %v", rcvd, msg)
    61  	}
    62  }