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

     1  package sync
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	core "github.com/libp2p/go-libp2p-core"
    10  	"github.com/libp2p/go-libp2p-core/network"
    11  	"github.com/prysmaticlabs/prysm/beacon-chain/p2p"
    12  	p2ptest "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
    13  	"github.com/prysmaticlabs/prysm/shared/testutil"
    14  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    15  )
    16  
    17  func TestContextWrite_NoWrites(t *testing.T) {
    18  	p1 := p2ptest.NewTestP2P(t)
    19  	nPeer := p2ptest.NewTestP2P(t)
    20  	p1.Connect(nPeer)
    21  
    22  	wg := new(sync.WaitGroup)
    23  	prID := p2p.RPCPingTopicV1
    24  	wg.Add(1)
    25  	nPeer.BHost.SetStreamHandler(core.ProtocolID(prID), func(stream network.Stream) {
    26  		wg.Done()
    27  		// no-op
    28  	})
    29  	strm, err := p1.BHost.NewStream(context.Background(), nPeer.PeerID(), p2p.RPCPingTopicV1)
    30  	assert.NoError(t, err)
    31  
    32  	// Nothing will be written to the stream
    33  	assert.NoError(t, writeContextToStream(strm, nil))
    34  	if testutil.WaitTimeout(wg, 1*time.Second) {
    35  		t.Fatal("Did not receive stream within 1 sec")
    36  	}
    37  }
    38  
    39  func TestContextRead_NoReads(t *testing.T) {
    40  	p1 := p2ptest.NewTestP2P(t)
    41  	nPeer := p2ptest.NewTestP2P(t)
    42  	p1.Connect(nPeer)
    43  
    44  	wg := new(sync.WaitGroup)
    45  	prID := p2p.RPCPingTopicV1
    46  	wg.Add(1)
    47  	wantedData := []byte{'A', 'B', 'C', 'D'}
    48  	nPeer.BHost.SetStreamHandler(core.ProtocolID(prID), func(stream network.Stream) {
    49  		// No Context will be read from it
    50  		dt, err := readContextFromStream(stream, nil)
    51  		assert.NoError(t, err)
    52  		assert.Equal(t, 0, len(dt))
    53  
    54  		// Ensure sent over data hasn't been modified.
    55  		buf := make([]byte, len(wantedData))
    56  		n, err := stream.Read(buf)
    57  		assert.NoError(t, err)
    58  		assert.Equal(t, len(wantedData), n)
    59  		assert.DeepEqual(t, wantedData, buf)
    60  
    61  		wg.Done()
    62  	})
    63  	strm, err := p1.BHost.NewStream(context.Background(), nPeer.PeerID(), p2p.RPCPingTopicV1)
    64  	assert.NoError(t, err)
    65  
    66  	n, err := strm.Write(wantedData)
    67  	assert.NoError(t, err)
    68  	assert.Equal(t, len(wantedData), n)
    69  	if testutil.WaitTimeout(wg, 1*time.Second) {
    70  		t.Fatal("Did not receive stream within 1 sec")
    71  	}
    72  }