fix(tests): stabilize unreliable session tests

session/session_test.go

@@ -35,13 +35,11 @@ func (fwm *fakeWantManager) CancelWants(ctx context.Context, cids []cid.Cid, pee
 type fakePeerManager struct {
 	lk                     sync.RWMutex
 	peers                  []peer.ID
-	findMorePeersRequested bool
+	findMorePeersRequested chan struct{}
 }
 
 func (fpm *fakePeerManager) FindMorePeers(context.Context, cid.Cid) {
-	fpm.lk.Lock()
-	fpm.findMorePeersRequested = true
-	fpm.lk.Unlock()
+	fpm.findMorePeersRequested <- struct{}{}
 }
 
 func (fpm *fakePeerManager) GetOptimizedPeers() []peer.ID {
@@ -161,7 +159,7 @@ func TestSessionFindMorePeers(t *testing.T) {
 	wantReqs := make(chan wantReq, 1)
 	cancelReqs := make(chan wantReq, 1)
 	fwm := &fakeWantManager{wantReqs, cancelReqs}
-	fpm := &fakePeerManager{}
+	fpm := &fakePeerManager{findMorePeersRequested: make(chan struct{})}
 	id := testutil.GenerateSessionID()
 	session := New(ctx, id, fwm, fpm)
 	session.SetBaseTickDelay(200 * time.Microsecond)
@@ -187,15 +185,8 @@ func TestSessionFindMorePeers(t *testing.T) {
 	<-wantReqs
 	<-cancelReqs
 
-	// wait long enough for a tick to occur
-	time.Sleep(20 * time.Millisecond)
-
-	// trigger to find providers should have happened
-	fpm.lk.Lock()
-	if fpm.findMorePeersRequested != true {
-		t.Fatal("should have attempted to find more peers but didn't")
-	}
-	fpm.lk.Unlock()
+	// wait for a request to get more peers to occur
+	<-fpm.findMorePeersRequested
 
 	// verify a broadcast was made
 	receivedWantReq := <-wantReqs

sessionpeermanager/sessionpeermanager_test.go

@@ -2,6 +2,7 @@ package sessionpeermanager
 
 import (
 	"context"
+	"sync"
 	"testing"
 	"time"
 
@@ -39,20 +40,27 @@ func (fpn *fakePeerNetwork) FindProvidersAsync(ctx context.Context, c cid.Cid, n
 
 type fakeConnManager struct {
 	taggedPeers []peer.ID
+	wait        sync.WaitGroup
 }
 
 func (fcm *fakeConnManager) TagPeer(p peer.ID, tag string, n int) {
+	fcm.wait.Add(1)
 	fcm.taggedPeers = append(fcm.taggedPeers, p)
 }
+
 func (fcm *fakeConnManager) UntagPeer(p peer.ID, tag string) {
+	defer fcm.wait.Done()
+
 	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 }
@@ -130,7 +138,8 @@ func TestUntaggingPeers(t *testing.T) {
 		t.Fatal("Peers were not tagged!")
 	}
 	<-ctx.Done()
-	time.Sleep(5 * time.Millisecond)
+	fcm.wait.Wait()
+
 	if len(fcm.taggedPeers) != 0 {
 		t.Fatal("Peers were not untagged!")
 	}