github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/node/internal/protocolPeerCache_test.go (about)

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