p2p/peer_set_test.go

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