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

     1  package sync
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/kevinms/leakybucket-go"
    10  	"github.com/libp2p/go-libp2p-core/network"
    11  	"github.com/libp2p/go-libp2p-core/protocol"
    12  	db "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
    13  	"github.com/prysmaticlabs/prysm/beacon-chain/p2p"
    14  	p2ptest "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
    15  	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
    16  	"github.com/prysmaticlabs/prysm/shared/interfaces"
    17  	"github.com/prysmaticlabs/prysm/shared/sszutil"
    18  	"github.com/prysmaticlabs/prysm/shared/testutil"
    19  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    20  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    21  )
    22  
    23  func TestMetaDataRPCHandler_ReceivesMetadata(t *testing.T) {
    24  	p1 := p2ptest.NewTestP2P(t)
    25  	p2 := p2ptest.NewTestP2P(t)
    26  	p1.Connect(p2)
    27  	assert.Equal(t, 1, len(p1.BHost.Network().Peers()), "Expected peers to be connected")
    28  	bitfield := [8]byte{'A', 'B'}
    29  	p1.LocalMetadata = interfaces.WrappedMetadataV0(&pb.MetaDataV0{
    30  		SeqNumber: 2,
    31  		Attnets:   bitfield[:],
    32  	})
    33  
    34  	// Set up a head state in the database with data we expect.
    35  	d := db.SetupDB(t)
    36  	r := &Service{
    37  		cfg: &Config{
    38  			DB:  d,
    39  			P2P: p1,
    40  		},
    41  		rateLimiter: newRateLimiter(p1),
    42  	}
    43  
    44  	// Setup streams
    45  	pcl := protocol.ID("/testing")
    46  	topic := string(pcl)
    47  	r.rateLimiter.limiterMap[topic] = leakybucket.NewCollector(1, 1, false)
    48  	var wg sync.WaitGroup
    49  	wg.Add(1)
    50  	p2.BHost.SetStreamHandler(pcl, func(stream network.Stream) {
    51  		defer wg.Done()
    52  		expectSuccess(t, stream)
    53  		out := new(pb.MetaDataV0)
    54  		assert.NoError(t, r.cfg.P2P.Encoding().DecodeWithMaxLength(stream, out))
    55  		assert.DeepEqual(t, p1.LocalMetadata.InnerObject(), out, "MetadataV0 unequal")
    56  	})
    57  	stream1, err := p1.BHost.NewStream(context.Background(), p2.BHost.ID(), pcl)
    58  	require.NoError(t, err)
    59  
    60  	assert.NoError(t, r.metaDataHandler(context.Background(), new(interface{}), stream1))
    61  
    62  	if testutil.WaitTimeout(&wg, 1*time.Second) {
    63  		t.Fatal("Did not receive stream within 1 sec")
    64  	}
    65  
    66  	conns := p1.BHost.Network().ConnsToPeer(p2.BHost.ID())
    67  	if len(conns) == 0 {
    68  		t.Error("Peer is disconnected despite receiving a valid ping")
    69  	}
    70  }
    71  
    72  func TestMetadataRPCHandler_SendsMetadata(t *testing.T) {
    73  	p1 := p2ptest.NewTestP2P(t)
    74  	p2 := p2ptest.NewTestP2P(t)
    75  	p1.Connect(p2)
    76  	assert.Equal(t, 1, len(p1.BHost.Network().Peers()), "Expected peers to be connected")
    77  	bitfield := [8]byte{'A', 'B'}
    78  	p2.LocalMetadata = interfaces.WrappedMetadataV0(&pb.MetaDataV0{
    79  		SeqNumber: 2,
    80  		Attnets:   bitfield[:],
    81  	})
    82  
    83  	// Set up a head state in the database with data we expect.
    84  	d := db.SetupDB(t)
    85  	r := &Service{
    86  		cfg: &Config{
    87  			DB:  d,
    88  			P2P: p1,
    89  		},
    90  		rateLimiter: newRateLimiter(p1),
    91  	}
    92  
    93  	r2 := &Service{
    94  		cfg: &Config{
    95  			DB:  d,
    96  			P2P: p2,
    97  		},
    98  		rateLimiter: newRateLimiter(p2),
    99  	}
   100  
   101  	// Setup streams
   102  	pcl := protocol.ID(p2p.RPCMetaDataTopicV1 + r.cfg.P2P.Encoding().ProtocolSuffix())
   103  	topic := string(pcl)
   104  	r.rateLimiter.limiterMap[topic] = leakybucket.NewCollector(1, 1, false)
   105  	r2.rateLimiter.limiterMap[topic] = leakybucket.NewCollector(1, 1, false)
   106  
   107  	var wg sync.WaitGroup
   108  	wg.Add(1)
   109  	p2.BHost.SetStreamHandler(pcl, func(stream network.Stream) {
   110  		defer wg.Done()
   111  		assert.NoError(t, r2.metaDataHandler(context.Background(), new(interface{}), stream))
   112  	})
   113  
   114  	metadata, err := r.sendMetaDataRequest(context.Background(), p2.BHost.ID())
   115  	assert.NoError(t, err)
   116  
   117  	if !sszutil.DeepEqual(metadata.InnerObject(), p2.LocalMetadata.InnerObject()) {
   118  		t.Fatalf("MetadataV0 unequal, received %v but wanted %v", metadata, p2.LocalMetadata)
   119  	}
   120  
   121  	if testutil.WaitTimeout(&wg, 1*time.Second) {
   122  		t.Fatal("Did not receive stream within 1 sec")
   123  	}
   124  
   125  	conns := p1.BHost.Network().ConnsToPeer(p2.BHost.ID())
   126  	if len(conns) == 0 {
   127  		t.Error("Peer is disconnected despite receiving a valid ping")
   128  	}
   129  }