From 78d4f3873f8b07c27f0fa16431e4d0ec488fef3b Mon Sep 17 00:00:00 2001
From: hannahhoward <hannah@hannahhoward.net>
Date: Tue, 18 Dec 2018 15:34:16 -0800
Subject: [PATCH 1/2] fix(tests): stabilize unreliable session tests

fix #43
---
 session/session_test.go                       | 17 ++++-------------
 sessionpeermanager/sessionpeermanager_test.go | 11 ++++++++++-
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/session/session_test.go b/session/session_test.go
index 1e6a8915..b00f8bd0 100644
--- a/session/session_test.go
+++ b/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)
@@ -188,14 +186,7 @@ func TestSessionFindMorePeers(t *testing.T) {
 	<-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()
+	<-fpm.findMorePeersRequested
 
 	// verify a broadcast was made
 	receivedWantReq := <-wantReqs
diff --git a/sessionpeermanager/sessionpeermanager_test.go b/sessionpeermanager/sessionpeermanager_test.go
index 821752a0..c26bf174 100644
--- a/sessionpeermanager/sessionpeermanager_test.go
+++ b/sessionpeermanager/sessionpeermanager_test.go
@@ -2,6 +2,7 @@ package sessionpeermanager
 
 import (
 	"context"
+	"sync"
 	"testing"
 	"time"
 
@@ -39,12 +40,17 @@ 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) {
+	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]
@@ -52,7 +58,9 @@ func (fcm *fakeConnManager) UntagPeer(p peer.ID, tag string) {
 			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!")
 	}

From 6b3042fe0ec4b5af0faa0db685a796e91bba2836 Mon Sep 17 00:00:00 2001
From: hannahhoward <hannah@hannahhoward.net>
Date: Tue, 18 Dec 2018 17:52:40 -0800
Subject: [PATCH 2/2] fix(tests): minor fix for waitgroup

---
 session/session_test.go                       | 2 +-
 sessionpeermanager/sessionpeermanager_test.go | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/session/session_test.go b/session/session_test.go
index b00f8bd0..8ae87cfd 100644
--- a/session/session_test.go
+++ b/session/session_test.go
@@ -185,7 +185,7 @@ func TestSessionFindMorePeers(t *testing.T) {
 	<-wantReqs
 	<-cancelReqs
 
-	// wait long enough for a tick to occur
+	// wait for a request to get more peers to occur
 	<-fpm.findMorePeersRequested
 
 	// verify a broadcast was made
diff --git a/sessionpeermanager/sessionpeermanager_test.go b/sessionpeermanager/sessionpeermanager_test.go
index c26bf174..f84b3d67 100644
--- a/sessionpeermanager/sessionpeermanager_test.go
+++ b/sessionpeermanager/sessionpeermanager_test.go
@@ -49,7 +49,7 @@ func (fcm *fakeConnManager) TagPeer(p peer.ID, tag string, n int) {
 }
 
 func (fcm *fakeConnManager) UntagPeer(p peer.ID, tag string) {
-	fcm.wait.Done()
+	defer fcm.wait.Done()
 
 	for i := 0; i < len(fcm.taggedPeers); i++ {
 		if fcm.taggedPeers[i] == p {