prioritized_peers_test.go

   1 package torrent
   2
   3 import (
   4         "net"
   5         "sort"
   6         "testing"
   7
   8         "github.com/google/btree"
   9         "github.com/stretchr/testify/assert"
  10 )
  11
  12 func TestPrioritizedPeers(t *testing.T) {
  13         pp := prioritizedPeers{
  14                 om: btree.New(3),
  15                 getPrio: func(p Peer) peerPriority {
  16                         return bep40PriorityIgnoreError(p.addr(), ipPort{IP: net.ParseIP("0.0.0.0")})
  17                 },
  18         }
  19         _, ok := pp.DeleteMin()
  20         assert.Panics(t, func() { pp.PopMax() })
  21         assert.False(t, ok)
  22         ps := []Peer{
  23                 {IP: net.ParseIP("1.2.3.4")},
  24                 {IP: net.ParseIP("1::2")},
  25                 {IP: net.ParseIP("")},
  26         }
  27         for i, p := range ps {
  28                 t.Logf("peer %d priority: %08x\n", i, pp.getPrio(p))
  29                 assert.False(t, pp.Add(p))
  30                 assert.True(t, pp.Add(p))
  31                 assert.Equal(t, i+1, pp.Len())
  32         }
  33         sort.Slice(ps, func(i, j int) bool {
  34                 return pp.getPrio(ps[i]) < pp.getPrio(ps[j])
  35         })
  36         pop := func(expected *Peer) {
  37                 if expected == nil {
  38                         assert.Panics(t, func() { pp.PopMax() })
  39                 } else {
  40                         assert.Equal(t, *expected, pp.PopMax())
  41                 }
  42         }
  43         min := func(expected *Peer) {
  44                 i, ok := pp.DeleteMin()
  45                 if expected == nil {
  46                         assert.False(t, ok)
  47                 } else {
  48                         assert.True(t, ok)
  49                         assert.Equal(t, *expected, i.p)
  50                 }
  51         }
  52         pop(&ps[2])
  53         min(&ps[0])
  54         pop(&ps[1])
  55         min(nil)
  56         pop(nil)
  57 }