From cefcda18c65863afa43bd1599975f41882177540 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sat, 14 Apr 2018 21:44:03 +1000
Subject: [PATCH] Add tests for prioritizedPeers

---
 prioritized_peers.go      | 10 +++++--
 prioritized_peers_test.go | 57 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)
 create mode 100644 prioritized_peers_test.go

diff --git a/prioritized_peers.go b/prioritized_peers.go
index 4ea45919..ca8b8736 100644
--- 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
index 00000000..df4c4b2d
--- /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)
+}
-- 
2.48.1