Add tests for prioritizedPeers

diff --git a/prioritized_peers.go b/prioritized_peers.go
index 4ea45919094a0fcfeb9e7bd1325463b6d8586c30..ca8b8736a626e007988ab650b105958dafec1dbc 100644 (file)
--- a/prioritized_peers.go
+++ b/prioritized_peers.go
@@ -34,8 +34,14 @@ func (me *prioritizedPeers) Add(p Peer) bool {
        return me.om.ReplaceOrInsert(prioritizedPeersItem{me.getPrio(p), p}) != nil
 }
 
-func (me *prioritizedPeers) DeleteMin() {
-       me.om.DeleteMin()
+func (me *prioritizedPeers) DeleteMin() (ret prioritizedPeersItem, ok bool) {
+       i := me.om.DeleteMin()
+       if i == nil {
+               return
+       }
+       ret = i.(prioritizedPeersItem)
+       ok = true
+       return
 }
 
 func (me *prioritizedPeers) PopMax() Peer {

diff --git a/prioritized_peers_test.go b/prioritized_peers_test.go
new file mode 100644 (file)
index 0000000..df4c4b2
--- /dev/null
+++ b/prioritized_peers_test.go
@@ -0,0 +1,57 @@
+package torrent
+
+import (
+       "net"
+       "sort"
+       "testing"
+
+       "github.com/google/btree"
+       "github.com/stretchr/testify/assert"
+)
+
+func TestPrioritizedPeers(t *testing.T) {
+       pp := prioritizedPeers{
+               om: btree.New(3),
+               getPrio: func(p Peer) peerPriority {
+                       return bep40PriorityIgnoreError(p.addr(), ipPort{IP: net.ParseIP("0.0.0.0")})
+               },
+       }
+       _, ok := pp.DeleteMin()
+       assert.Panics(t, func() { pp.PopMax() })
+       assert.False(t, ok)
+       ps := []Peer{
+               {IP: net.ParseIP("1.2.3.4")},
+               {IP: net.ParseIP("1::2")},
+               {IP: net.ParseIP("")},
+       }
+       for i, p := range ps {
+               t.Logf("peer %d priority: %08x\n", i, pp.getPrio(p))
+               assert.False(t, pp.Add(p))
+               assert.True(t, pp.Add(p))
+               assert.Equal(t, i+1, pp.Len())
+       }
+       sort.Slice(ps, func(i, j int) bool {
+               return pp.getPrio(ps[i]) < pp.getPrio(ps[j])
+       })
+       pop := func(expected *Peer) {
+               if expected == nil {
+                       assert.Panics(t, func() { pp.PopMax() })
+               } else {
+                       assert.Equal(t, *expected, pp.PopMax())
+               }
+       }
+       min := func(expected *Peer) {
+               i, ok := pp.DeleteMin()
+               if expected == nil {
+                       assert.False(t, ok)
+               } else {
+                       assert.True(t, ok)
+                       assert.Equal(t, *expected, i.p)
+               }
+       }
+       pop(&ps[2])
+       min(&ps[0])
+       pop(&ps[1])
+       min(nil)
+       pop(nil)
+}