Adds a test for PeerTracker (#576)

peer/peer_tracker_test.go

@@ -0,0 +1,88 @@
+// (c) 2019-2022, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package peer
+
+import (
+	"testing"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPeerTracker(t *testing.T) {
+	require := require.New(t)
+	p := NewPeerTracker()
+
+	// Connect some peers
+	numExtraPeers := 10
+	numPeers := desiredMinResponsivePeers + numExtraPeers
+	peerIDs := make([]ids.NodeID, numPeers)
+
+	for i := range peerIDs {
+		peerIDs[i] = ids.GenerateTestNodeID()
+		p.Connected(peerIDs[i], defaultPeerVersion)
+	}
+
+	responsivePeers := make(map[ids.NodeID]bool)
+
+	// Expect requests to go to new peers until we have desiredMinResponsivePeers responsive peers.
+	for i := 0; i < desiredMinResponsivePeers+numExtraPeers/2; i++ {
+		peer, ok := p.GetAnyPeer(nil)
+		require.True(ok)
+		require.NotNil(peer)
+
+		_, exists := responsivePeers[peer]
+		require.Falsef(exists, "expected connecting to a new peer, but got the same peer twice: peer %s iteration %d", peer, i)
+		responsivePeers[peer] = true
+
+		p.TrackPeer(peer) // mark the peer as having a message sent to it
+	}
+
+	// Mark some peers as responsive and others as not responsive
+	i := 0
+	for peer := range responsivePeers {
+		if i < desiredMinResponsivePeers {
+			p.TrackBandwidth(peer, 10)
+		} else {
+			responsivePeers[peer] = false // remember which peers were not responsive
+			p.TrackBandwidth(peer, 0)
+		}
+		i++
+	}
+
+	// Expect requests to go to responsive or new peers, so long as they are available
+	numRequests := 50
+	for i := 0; i < numRequests; i++ {
+		peer, ok := p.GetAnyPeer(nil)
+		require.True(ok)
+		require.NotNil(peer)
+
+		responsive, ok := responsivePeers[peer]
+		if ok {
+			require.Truef(responsive, "expected connecting to a responsive peer, but got a peer that was not responsive: peer %s iteration %d", peer, i)
+			p.TrackBandwidth(peer, 10)
+		} else {
+			responsivePeers[peer] = false // remember that we connected to this peer
+			p.TrackPeer(peer)             // mark the peer as having a message sent to it
+			p.TrackBandwidth(peer, 0)     // mark the peer as non-responsive
+		}
+	}
+
+	// Disconnect from peers that were previously responsive and ones we didn't connect to yet.
+	for _, peer := range peerIDs {
+		responsive, ok := responsivePeers[peer]
+		if ok && responsive || !ok {
+			p.Disconnected(peer)
+		}
+	}
+
+	// Requests should fall back on non-responsive peers when no other choice is left
+	peer, ok := p.GetAnyPeer(nil)
+	require.True(ok)
+	require.NotNil(peer)
+
+	responsive, ok := responsivePeers[peer]
+	require.True(ok)
+	require.Falsef(responsive, "expected connecting to a non-responsive peer, but got a peer that was responsive: peer %s", peer)
+}