github.com/iotexproject/iotex-core@v1.14.1-rc1/p2p/agent_test.go (about)

     1  // Copyright (c) 2019 IoTeX Foundation
     2  // This source code is provided 'as is' and no warranties are given as to title or non-infringement, merchantability
     3  // or fitness for purpose and, to the extent permitted by law, all liability for your use of the code is disclaimed.
     4  // This source code is governed by Apache License 2.0 that can be found in the LICENSE file.
     5  
     6  package p2p
     7  
     8  import (
     9  	"context"
    10  	"strconv"
    11  	"sync"
    12  	"testing"
    13  	"time"
    14  
    15  	"github.com/libp2p/go-libp2p-core/peer"
    16  	"github.com/stretchr/testify/require"
    17  	"google.golang.org/protobuf/proto"
    18  
    19  	"github.com/iotexproject/go-p2p"
    20  	"github.com/iotexproject/go-pkgs/hash"
    21  	"github.com/iotexproject/iotex-proto/golang/testingpb"
    22  
    23  	"github.com/iotexproject/iotex-core/testutil"
    24  )
    25  
    26  func TestDummyAgent(t *testing.T) {
    27  	require := require.New(t)
    28  	a := NewDummyAgent()
    29  	require.NoError(a.Start(nil))
    30  	require.NoError(a.Stop(nil))
    31  	require.NoError(a.BroadcastOutbound(nil, nil))
    32  	require.NoError(a.UnicastOutbound(nil, peer.AddrInfo{}, nil))
    33  	info, err := a.Info()
    34  	require.Equal(peer.AddrInfo{}, info)
    35  	require.NoError(err)
    36  	addrs, err := a.Self()
    37  	require.Nil(addrs)
    38  	require.NoError(err)
    39  	neighbors, err := a.ConnectedPeers()
    40  	require.Nil(neighbors)
    41  	require.NoError(err)
    42  }
    43  
    44  func TestBroadcast(t *testing.T) {
    45  	r := require.New(t)
    46  
    47  	ctx := context.Background()
    48  	n := 10
    49  	agents := make([]Agent, 0)
    50  	defer func() {
    51  		var err error
    52  		for _, agent := range agents {
    53  			err = agent.Stop(ctx)
    54  		}
    55  		r.NoError(err)
    56  	}()
    57  	counts := make(map[uint8]int)
    58  	var mutex sync.RWMutex
    59  	b := func(_ context.Context, _ uint32, _ string, msg proto.Message) {
    60  		mutex.Lock()
    61  		defer mutex.Unlock()
    62  		testMsg, ok := msg.(*testingpb.TestPayload)
    63  		r.True(ok)
    64  		idx := testMsg.MsgBody[0]
    65  		if _, ok = counts[idx]; ok {
    66  			counts[idx]++
    67  		} else {
    68  			counts[idx] = 1
    69  		}
    70  	}
    71  	u := func(_ context.Context, _ uint32, _ peer.AddrInfo, _ proto.Message) {}
    72  
    73  	bootnodePort := testutil.RandomPort()
    74  	bootnode, err := p2p.NewHost(context.Background(), p2p.DHTProtocolID(1), p2p.Port(bootnodePort), p2p.SecureIO(), p2p.MasterKey("bootnode"))
    75  	r.NoError(err)
    76  	bootnodeAddr := bootnode.Addresses()
    77  
    78  	for i := 0; i < n; i++ {
    79  		port := bootnodePort + i + 1
    80  		agent := NewAgent(Config{
    81  			Host:              "127.0.0.1",
    82  			Port:              port,
    83  			BootstrapNodes:    []string{bootnodeAddr[0].String()},
    84  			ReconnectInterval: 150 * time.Second,
    85  			MaxMessageSize:    p2p.DefaultConfig.MaxMessageSize,
    86  		}, 1, hash.ZeroHash256, b, u)
    87  		agent.Start(ctx)
    88  		agents = append(agents, agent)
    89  	}
    90  
    91  	for i := 0; i < n; i++ {
    92  		r.NoError(agents[i].BroadcastOutbound(ctx, &testingpb.TestPayload{
    93  			MsgBody: []byte{uint8(i)},
    94  		}))
    95  		r.NoError(testutil.WaitUntil(100*time.Millisecond, 20*time.Second, func() (bool, error) {
    96  			mutex.RLock()
    97  			defer mutex.RUnlock()
    98  			// Broadcast message will be skipped by the source node
    99  			return counts[uint8(i)] == n-1, nil
   100  		}))
   101  	}
   102  }
   103  
   104  func TestUnicast(t *testing.T) {
   105  	r := require.New(t)
   106  
   107  	ctx := context.Background()
   108  	n := 10
   109  	agents := make([]Agent, 0)
   110  	defer func() {
   111  		var err error
   112  		for _, agent := range agents {
   113  			err = agent.Stop(ctx)
   114  		}
   115  		r.NoError(err)
   116  	}()
   117  	counts := make(map[uint8]int)
   118  	var src string
   119  	var mutex sync.RWMutex
   120  	b := func(_ context.Context, _ uint32, _ string, _ proto.Message) {}
   121  	u := func(_ context.Context, _ uint32, peer peer.AddrInfo, msg proto.Message) {
   122  		mutex.Lock()
   123  		defer mutex.Unlock()
   124  		testMsg, ok := msg.(*testingpb.TestPayload)
   125  		r.True(ok)
   126  		idx := testMsg.MsgBody[0]
   127  		if _, ok = counts[idx]; ok {
   128  			counts[idx]++
   129  		} else {
   130  			counts[idx] = 1
   131  		}
   132  		src = peer.ID.Pretty()
   133  	}
   134  
   135  	bootnode, err := p2p.NewHost(context.Background(), p2p.DHTProtocolID(2), p2p.Port(testutil.RandomPort()), p2p.SecureIO(), p2p.MasterKey("bootnode"))
   136  	r.NoError(err)
   137  	addrs := bootnode.Addresses()
   138  	for i := 0; i < n; i++ {
   139  		agent := NewAgent(Config{
   140  			Host:              "127.0.0.1",
   141  			Port:              testutil.RandomPort(),
   142  			BootstrapNodes:    []string{addrs[0].String()},
   143  			ReconnectInterval: 150 * time.Second,
   144  			MasterKey:         strconv.Itoa(i),
   145  		}, 2, hash.ZeroHash256, b, u)
   146  		r.NoError(agent.Start(ctx))
   147  		agents = append(agents, agent)
   148  	}
   149  
   150  	for i := 0; i < n; i++ {
   151  		neighbors, err := agents[i].ConnectedPeers()
   152  		r.NoError(err)
   153  		r.True(len(neighbors) >= n/3)
   154  		for _, neighbor := range neighbors {
   155  			r.NoError(agents[i].UnicastOutbound(ctx, neighbor, &testingpb.TestPayload{
   156  				MsgBody: []byte{uint8(i)},
   157  			}))
   158  		}
   159  		r.NoError(testutil.WaitUntil(100*time.Millisecond, 20*time.Second, func() (bool, error) {
   160  			mutex.RLock()
   161  			defer mutex.RUnlock()
   162  			info, err := agents[i].Info()
   163  			r.NoError(err)
   164  			return counts[uint8(i)] == len(neighbors) && src == info.ID.Pretty(), nil
   165  		}))
   166  	}
   167  }