github.com/koko1123/flow-go-1@v0.29.6/network/p2p/p2pnode/protocolPeerCache_test.go (about)

     1  package p2pnode_test
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/libp2p/go-libp2p/core/host"
     9  	"github.com/libp2p/go-libp2p/core/network"
    10  	"github.com/libp2p/go-libp2p/core/protocol"
    11  	"github.com/rs/zerolog"
    12  	"github.com/stretchr/testify/assert"
    13  	"github.com/stretchr/testify/require"
    14  
    15  	"github.com/koko1123/flow-go-1/network/p2p/p2pbuilder"
    16  	"github.com/koko1123/flow-go-1/network/p2p/p2pnode"
    17  
    18  	fcrypto "github.com/onflow/flow-go/crypto"
    19  
    20  	"github.com/koko1123/flow-go-1/utils/unittest"
    21  )
    22  
    23  func TestProtocolPeerCache(t *testing.T) {
    24  	ctx, cancel := context.WithCancel(context.Background())
    25  	defer cancel()
    26  
    27  	// create three hosts, and a pcache for the first
    28  	h1, err := p2pbuilder.DefaultLibP2PHost(unittest.DefaultAddress, unittest.KeyFixture(fcrypto.ECDSASecp256k1))
    29  	require.NoError(t, err)
    30  	pcache, err := p2pnode.NewProtocolPeerCache(zerolog.Nop(), h1)
    31  	require.NoError(t, err)
    32  	h2, err := p2pbuilder.DefaultLibP2PHost(unittest.DefaultAddress, unittest.KeyFixture(fcrypto.ECDSASecp256k1))
    33  	require.NoError(t, err)
    34  	h3, err := p2pbuilder.DefaultLibP2PHost(unittest.DefaultAddress, unittest.KeyFixture(fcrypto.ECDSASecp256k1))
    35  	require.NoError(t, err)
    36  
    37  	// register each host on a separate protocol
    38  	p1 := protocol.ID("p1")
    39  	p2 := protocol.ID("p2")
    40  	p3 := protocol.ID("p3")
    41  	noopHandler := func(s network.Stream) {}
    42  	h1.SetStreamHandler(p1, noopHandler)
    43  	h2.SetStreamHandler(p2, noopHandler)
    44  	h3.SetStreamHandler(p3, noopHandler)
    45  
    46  	// connect the hosts to each other
    47  	require.NoError(t, h1.Connect(ctx, *host.InfoFromHost(h2)))
    48  	require.NoError(t, h1.Connect(ctx, *host.InfoFromHost(h3)))
    49  	require.NoError(t, h2.Connect(ctx, *host.InfoFromHost(h3)))
    50  
    51  	// check that h1's pcache reflects the protocols supported by h2 and h3
    52  	assert.Eventually(t, func() bool {
    53  		peers2 := pcache.GetPeers(p2)
    54  		peers3 := pcache.GetPeers(p3)
    55  		_, ok2 := peers2[h2.ID()]
    56  		_, ok3 := peers3[h3.ID()]
    57  		return len(peers2) == 1 && len(peers3) == 1 && ok2 && ok3
    58  	}, 3*time.Second, 50*time.Millisecond)
    59  
    60  	// remove h2's support for p2
    61  	h2.RemoveStreamHandler(p2)
    62  
    63  	// check that h1's pcache reflects the change
    64  	assert.Eventually(t, func() bool {
    65  		return len(pcache.GetPeers(p2)) == 0
    66  	}, 3*time.Second, 50*time.Millisecond)
    67  
    68  	// add support for p4 on h2 and h3
    69  	p4 := protocol.ID("p4")
    70  	h2.SetStreamHandler(p4, noopHandler)
    71  	h3.SetStreamHandler(p4, noopHandler)
    72  
    73  	// check that h1's pcache reflects the change
    74  	assert.Eventually(t, func() bool {
    75  		peers4 := pcache.GetPeers(p4)
    76  		_, ok2 := peers4[h2.ID()]
    77  		_, ok3 := peers4[h3.ID()]
    78  		return len(peers4) == 2 && ok2 && ok3
    79  	}, 3*time.Second, 50*time.Millisecond)
    80  }