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

     1  package connection_test
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  	"testing"
     7  
     8  	"github.com/libp2p/go-libp2p/core/peer"
     9  	"github.com/rs/zerolog"
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/koko1123/flow-go-1/network/internal/p2pfixtures"
    13  	"github.com/koko1123/flow-go-1/network/p2p/connection"
    14  	"github.com/koko1123/flow-go-1/network/p2p/utils"
    15  
    16  	"github.com/koko1123/flow-go-1/module/metrics"
    17  	"github.com/koko1123/flow-go-1/utils/unittest"
    18  )
    19  
    20  const (
    21  	protectF     = "protect"
    22  	unprotectF   = "unprotect"
    23  	isProtectedF = "isprotected"
    24  )
    25  
    26  type fun struct {
    27  	funName     string
    28  	expectation bool
    29  }
    30  
    31  var protect = fun{
    32  	protectF,
    33  	false,
    34  }
    35  var unprotect = fun{
    36  	unprotectF,
    37  	false,
    38  }
    39  var isProtected = fun{
    40  	isProtectedF,
    41  	true,
    42  }
    43  var isNotProtected = fun{
    44  	isProtectedF,
    45  	false,
    46  }
    47  
    48  // TestConnectionManagerProtection tests that multiple protect and unprotect calls result in the correct IsProtected
    49  // status for a peer ID
    50  func TestConnectionManagerProtection(t *testing.T) {
    51  
    52  	log := zerolog.New(os.Stderr).Level(zerolog.ErrorLevel)
    53  	noopMetrics := metrics.NewNoopCollector()
    54  	connManager := connection.NewConnManager(log, noopMetrics)
    55  
    56  	testCases := [][]fun{
    57  		// single stream created on a connection
    58  		{protect, isProtected, unprotect, isNotProtected},
    59  		// two streams created on a connection at the same time
    60  		{protect, protect, unprotect, isNotProtected, unprotect, isNotProtected},
    61  		// two streams created on a connection one after another
    62  		{protect, unprotect, isNotProtected, protect, unprotect, isNotProtected},
    63  	}
    64  
    65  	for _, testCase := range testCases {
    66  		testSequence(t, testCase, connManager)
    67  	}
    68  }
    69  
    70  func testSequence(t *testing.T, sequence []fun, connMgr *connection.ConnManager) {
    71  	pID := generatePeerInfo(t)
    72  	for _, s := range sequence {
    73  		switch s.funName {
    74  		case protectF:
    75  			connMgr.Protect(pID, "global")
    76  		case unprotectF:
    77  			connMgr.Unprotect(pID, "global")
    78  		case isProtectedF:
    79  			require.Equal(t, connMgr.IsProtected(pID, ""), s.expectation, fmt.Sprintf("failed sequence: %v", sequence))
    80  		}
    81  	}
    82  }
    83  
    84  func generatePeerInfo(t *testing.T) peer.ID {
    85  	key := p2pfixtures.NetworkingKeyFixtures(t)
    86  	identity := unittest.IdentityFixture(unittest.WithNetworkingKey(key.PublicKey()), unittest.WithAddress("1.1.1.1:0"))
    87  	pInfo, err := utils.PeerAddressInfo(*identity)
    88  	require.NoError(t, err)
    89  	return pInfo.ID
    90  }