FedP2P/prioritized_peers_test.go

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)
}