p2p/peer_set_test.go

   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 }