test(sessionpeermanager): Add unit test

Add unit test for sessionpeermanger and comment exported methods

sessionpeermanager/sessionpeermanager.go

@@ -9,11 +9,14 @@ import (
 	peer "github.com/libp2p/go-libp2p-peer"
 )
 
+// PeerNetwork is an interface for finding providers and managing connections
 type PeerNetwork interface {
 	ConnectionManager() ifconnmgr.ConnManager
 	FindProvidersAsync(context.Context, cid.Cid, int) <-chan peer.ID
 }
 
+// SessionPeerManager tracks and manages peers for a session, and provides
+// the best ones to the session
 type SessionPeerManager struct {
 	ctx     context.Context
 	network PeerNetwork
@@ -27,6 +30,7 @@ type SessionPeerManager struct {
 	activePeersArr []peer.ID
 }
 
+// New creates a new SessionPeerManager
 func New(ctx context.Context, id uint64, network PeerNetwork) *SessionPeerManager {
 	spm := &SessionPeerManager{
 		ctx:         ctx,
@@ -42,7 +46,10 @@ func New(ctx context.Context, id uint64, network PeerNetwork) *SessionPeerManage
 	return spm
 }
 
+// RecordPeerResponse records that a peer received a block, and adds to it
+// the list of peers if it wasn't already added
 func (spm *SessionPeerManager) RecordPeerResponse(p peer.ID, k cid.Cid) {
+
 	// at the moment, we're just adding peers here
 	// in the future, we'll actually use this to record metrics
 	select {
@@ -51,11 +58,13 @@ func (spm *SessionPeerManager) RecordPeerResponse(p peer.ID, k cid.Cid) {
 	}
 }
 
+// RecordPeerRequests records that a given set of peers requested the given cids
 func (spm *SessionPeerManager) RecordPeerRequests(p []peer.ID, ks []cid.Cid) {
 	// at the moment, we're not doing anything here
 	// soon we'll use this to track latency by peer
 }
 
+// GetOptimizedPeers returns the best peers available for a session
 func (spm *SessionPeerManager) GetOptimizedPeers() []peer.ID {
 	// right now this just returns all peers, but soon we might return peers
 	// ordered by optimization, or only a subset
@@ -74,6 +83,8 @@ func (spm *SessionPeerManager) GetOptimizedPeers() []peer.ID {
 	}
 }
 
+// FindMorePeers attempts to find more peers for a session by searching for
+// providers for the given Cid
 func (spm *SessionPeerManager) FindMorePeers(ctx context.Context, c cid.Cid) {
 	go func(k cid.Cid) {
 		// TODO: have a task queue setup for this to:

sessionpeermanager/sessionpeermanager_test.go

+package sessionpeermanager
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/ipfs/go-bitswap/testutil"
+
+	cid "github.com/ipfs/go-cid"
+	ifconnmgr "github.com/libp2p/go-libp2p-interface-connmgr"
+	inet "github.com/libp2p/go-libp2p-net"
+	peer "github.com/libp2p/go-libp2p-peer"
+)
+
+type fakePeerNetwork struct {
+	peers       []peer.ID
+	connManager ifconnmgr.ConnManager
+}
+
+func (fpn *fakePeerNetwork) ConnectionManager() ifconnmgr.ConnManager {
+	return fpn.connManager
+}
+
+func (fpn *fakePeerNetwork) FindProvidersAsync(ctx context.Context, c cid.Cid, num int) <-chan peer.ID {
+	peerCh := make(chan peer.ID)
+	go func() {
+		defer close(peerCh)
+		for _, p := range fpn.peers {
+			select {
+			case peerCh <- p:
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+	return peerCh
+}
+
+type fakeConnManager struct {
+	taggedPeers []peer.ID
+}
+
+func (fcm *fakeConnManager) TagPeer(p peer.ID, tag string, n int) {
+	fcm.taggedPeers = append(fcm.taggedPeers, p)
+}
+func (fcm *fakeConnManager) UntagPeer(p peer.ID, tag string) {
+	for i := 0; i < len(fcm.taggedPeers); i++ {
+		if fcm.taggedPeers[i] == p {
+			fcm.taggedPeers[i] = fcm.taggedPeers[len(fcm.taggedPeers)-1]
+			fcm.taggedPeers = fcm.taggedPeers[:len(fcm.taggedPeers)-1]
+			return
+		}
+	}
+}
+func (*fakeConnManager) GetTagInfo(p peer.ID) *ifconnmgr.TagInfo { return nil }
+func (*fakeConnManager) TrimOpenConns(ctx context.Context)       {}
+func (*fakeConnManager) Notifee() inet.Notifiee                  { return nil }
+
+func TestFindingMorePeers(t *testing.T) {
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	peers := testutil.GeneratePeers(5)
+	fcm := &fakeConnManager{}
+	fpn := &fakePeerNetwork{peers, fcm}
+	c := testutil.GenerateCids(1)[0]
+	id := testutil.GenerateSessionID()
+
+	sessionPeerManager := New(ctx, id, fpn)
+
+	findCtx, findCancel := context.WithTimeout(ctx, 10*time.Millisecond)
+	defer findCancel()
+	sessionPeerManager.FindMorePeers(ctx, c)
+	<-findCtx.Done()
+	sessionPeers := sessionPeerManager.GetOptimizedPeers()
+	if len(sessionPeers) != len(peers) {
+		t.Fatal("incorrect number of peers found")
+	}
+	for _, p := range sessionPeers {
+		if !testutil.ContainsPeer(peers, p) {
+			t.Fatal("incorrect peer found through finding providers")
+		}
+	}
+	if len(fcm.taggedPeers) != len(peers) {
+		t.Fatal("Peers were not tagged!")
+	}
+}
+
+func TestRecordingReceivedBlocks(t *testing.T) {
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	p := testutil.GeneratePeers(1)[0]
+	fcm := &fakeConnManager{}
+	fpn := &fakePeerNetwork{nil, fcm}
+	c := testutil.GenerateCids(1)[0]
+	id := testutil.GenerateSessionID()
+
+	sessionPeerManager := New(ctx, id, fpn)
+	sessionPeerManager.RecordPeerResponse(p, c)
+	time.Sleep(10 * time.Millisecond)
+	sessionPeers := sessionPeerManager.GetOptimizedPeers()
+	if len(sessionPeers) != 1 {
+		t.Fatal("did not add peer on receive")
+	}
+	if sessionPeers[0] != p {
+		t.Fatal("incorrect peer added on receive")
+	}
+	if len(fcm.taggedPeers) != 1 {
+		t.Fatal("Peers was not tagged!")
+	}
+}
+
+func TestUntaggingPeers(t *testing.T) {
+	ctx := context.Background()
+	ctx, cancel := context.WithTimeout(ctx, 10*time.Millisecond)
+	defer cancel()
+	peers := testutil.GeneratePeers(5)
+	fcm := &fakeConnManager{}
+	fpn := &fakePeerNetwork{peers, fcm}
+	c := testutil.GenerateCids(1)[0]
+	id := testutil.GenerateSessionID()
+
+	sessionPeerManager := New(ctx, id, fpn)
+
+	sessionPeerManager.FindMorePeers(ctx, c)
+	time.Sleep(5 * time.Millisecond)
+	if len(fcm.taggedPeers) != len(peers) {
+		t.Fatal("Peers were not tagged!")
+	}
+	<-ctx.Done()
+	if len(fcm.taggedPeers) != 0 {
+		t.Fatal("Peers were not untagged!")
+	}
+}