github.com/Bytom/bytom@v1.1.2-0.20210127130405-ae40204c0b09/p2p/peer_set_test.go (about)

     1  // +build !network
     2  
     3  package p2p
     4  
     5  import (
     6  	"math/rand"
     7  	"testing"
     8  
     9  	cmn "github.com/tendermint/tmlibs/common"
    10  )
    11  
    12  // Returns an empty dummy peer
    13  func randPeer() *Peer {
    14  	return &Peer{
    15  		Key: cmn.RandStr(12),
    16  		NodeInfo: &NodeInfo{
    17  			RemoteAddr: cmn.Fmt("%v.%v.%v.%v:46656", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256),
    18  			ListenAddr: cmn.Fmt("%v.%v.%v.%v:46656", rand.Int()%256, rand.Int()%256, rand.Int()%256, rand.Int()%256),
    19  		},
    20  	}
    21  }
    22  
    23  func TestAddRemoveOne(t *testing.T) {
    24  	peerSet := NewPeerSet()
    25  
    26  	peer := randPeer()
    27  	err := peerSet.Add(peer)
    28  	if err != nil {
    29  		t.Errorf("Failed to add new peer")
    30  	}
    31  	if peerSet.Size() != 1 {
    32  		t.Errorf("Failed to add new peer and increment size")
    33  	}
    34  
    35  	peerSet.Remove(peer)
    36  	if peerSet.Has(peer.Key) {
    37  		t.Errorf("Failed to remove peer")
    38  	}
    39  	if peerSet.Size() != 0 {
    40  		t.Errorf("Failed to remove peer and decrement size")
    41  	}
    42  }
    43  
    44  func TestAddRemoveMany(t *testing.T) {
    45  	peerSet := NewPeerSet()
    46  
    47  	peers := []*Peer{}
    48  	N := 100
    49  	for i := 0; i < N; i++ {
    50  		peer := randPeer()
    51  		if err := peerSet.Add(peer); err != nil {
    52  			t.Errorf("Failed to add new peer")
    53  		}
    54  		if peerSet.Size() != i+1 {
    55  			t.Errorf("Failed to add new peer and increment size")
    56  		}
    57  		peers = append(peers, peer)
    58  	}
    59  
    60  	for i, peer := range peers {
    61  		peerSet.Remove(peer)
    62  		if peerSet.Has(peer.Key) {
    63  			t.Errorf("Failed to remove peer")
    64  		}
    65  		if peerSet.Size() != len(peers)-i-1 {
    66  			t.Errorf("Failed to remove peer and decrement size")
    67  		}
    68  	}
    69  }