Piece priority tests

diff --git a/connection_test.go b/connection_test.go
index 22b76f01c2ecd38ea6733837cc47d7a6c8ec1128..55181a9ca609bf2402f71ad405d00db109f7c1a1 100644 (file)
--- a/connection_test.go
+++ b/connection_test.go
@@ -4,6 +4,10 @@ import (
        "testing"
        "time"
 
+       "github.com/bradfitz/iter"
+
+       "bitbucket.org/anacrolix/go.torrent/internal/pieceordering"
+
        "bitbucket.org/anacrolix/go.torrent/peer_protocol"
 )
 
@@ -48,3 +52,44 @@ func TestCancelRequestOptimized(t *testing.T) {
                }
        }
 }
+
+func testRequestOrder(expected []int, ro *pieceordering.Instance, t *testing.T) {
+       e := ro.First()
+       for _, i := range expected {
+               if i != e.Piece() {
+                       t.FailNow()
+               }
+               e = e.Next()
+       }
+       if e != nil {
+               t.FailNow()
+       }
+}
+
+// Tests the request ordering based on a connections priorities.
+func TestPieceRequestOrder(t *testing.T) {
+       c := connection{
+               pieceRequestOrder: pieceordering.New(),
+               piecePriorities:   []int{1, 4, 0, 3, 2},
+       }
+       testRequestOrder(nil, c.pieceRequestOrder, t)
+       c.pendPiece(2, piecePriorityNone)
+       testRequestOrder(nil, c.pieceRequestOrder, t)
+       c.pendPiece(1, piecePriorityNormal)
+       c.pendPiece(2, piecePriorityNormal)
+       testRequestOrder([]int{2, 1}, c.pieceRequestOrder, t)
+       c.pendPiece(0, piecePriorityNormal)
+       testRequestOrder([]int{2, 0, 1}, c.pieceRequestOrder, t)
+       c.pendPiece(1, piecePriorityReadahead)
+       testRequestOrder([]int{1, 2, 0}, c.pieceRequestOrder, t)
+       c.pendPiece(4, piecePriorityNow)
+       testRequestOrder([]int{4, 1, 2, 0}, c.pieceRequestOrder, t)
+       c.pendPiece(2, piecePriorityReadahead)
+       // N(4), R(1, 2), N(0)
+       testRequestOrder([]int{4, 1, 2, 0}, c.pieceRequestOrder, t)
+       // Note this intentially sets to None a piece that's not in the order.
+       for i := range iter.N(5) {
+               c.pendPiece(i, piecePriorityNone)
+       }
+       testRequestOrder(nil, c.pieceRequestOrder, t)
+}