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

     1  package sync
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	libp2pcore "github.com/libp2p/go-libp2p-core"
    10  	"github.com/libp2p/go-libp2p-core/network"
    11  	"github.com/prysmaticlabs/prysm/beacon-chain/core/state"
    12  	prysmP2P "github.com/prysmaticlabs/prysm/beacon-chain/p2p"
    13  	"github.com/prysmaticlabs/prysm/beacon-chain/p2p/encoder"
    14  	p2ptest "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
    15  	p2ppb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
    16  	"github.com/prysmaticlabs/prysm/shared/testutil"
    17  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    18  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    19  )
    20  
    21  func init() {
    22  	state.SkipSlotCache.Disable()
    23  }
    24  
    25  // expectSuccess status code from a stream in regular sync.
    26  func expectSuccess(t *testing.T, stream network.Stream) {
    27  	code, errMsg, err := ReadStatusCode(stream, &encoder.SszNetworkEncoder{})
    28  	require.NoError(t, err)
    29  	require.Equal(t, uint8(0), code, "Received non-zero response code")
    30  	require.Equal(t, "", errMsg, "Received error message from stream")
    31  }
    32  
    33  // expectSuccess status code from a stream in regular sync.
    34  func expectFailure(t *testing.T, expectedCode uint8, expectedErrorMsg string, stream network.Stream) {
    35  	code, errMsg, err := ReadStatusCode(stream, &encoder.SszNetworkEncoder{})
    36  	require.NoError(t, err)
    37  	require.NotEqual(t, uint8(0), code, "Expected request to fail but got a 0 response code")
    38  	require.Equal(t, expectedCode, code, "Received incorrect response code")
    39  	require.Equal(t, expectedErrorMsg, errMsg)
    40  }
    41  
    42  // expectResetStream status code from a stream in regular sync.
    43  func expectResetStream(t *testing.T, stream network.Stream) {
    44  	expectedErr := "stream reset"
    45  	_, _, err := ReadStatusCode(stream, &encoder.SszNetworkEncoder{})
    46  	require.ErrorContains(t, expectedErr, err)
    47  }
    48  
    49  func TestRegisterRPC_ReceivesValidMessage(t *testing.T) {
    50  	p2p := p2ptest.NewTestP2P(t)
    51  	r := &Service{
    52  		ctx:         context.Background(),
    53  		cfg:         &Config{P2P: p2p},
    54  		rateLimiter: newRateLimiter(p2p),
    55  	}
    56  
    57  	var wg sync.WaitGroup
    58  	wg.Add(1)
    59  	topic := "/testing/foobar/1"
    60  	handler := func(ctx context.Context, msg interface{}, stream libp2pcore.Stream) error {
    61  		m, ok := msg.(*p2ppb.Fork)
    62  		if !ok {
    63  			t.Error("Object is not of type *pb.TestSimpleMessage")
    64  		}
    65  		assert.DeepEqual(t, []byte("fooo"), m.CurrentVersion, "Unexpected incoming message")
    66  		wg.Done()
    67  
    68  		return nil
    69  	}
    70  	prysmP2P.RPCTopicMappings[topic] = new(p2ppb.Fork)
    71  	// Cleanup Topic mappings
    72  	defer func() {
    73  		delete(prysmP2P.RPCTopicMappings, topic)
    74  	}()
    75  	r.registerRPC(topic, handler)
    76  
    77  	p2p.ReceiveRPC(topic, &p2ppb.Fork{CurrentVersion: []byte("fooo"), PreviousVersion: []byte("barr")})
    78  
    79  	if testutil.WaitTimeout(&wg, time.Second) {
    80  		t.Fatal("Did not receive RPC in 1 second")
    81  	}
    82  }