From 112866c35a6fa6e2f4a128e464a86d55110c0315 Mon Sep 17 00:00:00 2001
From: Steven Allen <steven@stebalien.com>
Date: Thu, 27 Jul 2017 14:35:29 -0700
Subject: [PATCH] tests: update test cases to act like go-peerstream 1.7.0

---
 p2p/net/mock/mock_notif_test.go | 19 ++++++++++---------
 p2p/net/mock/mock_peernet.go    | 14 ++++++++++----
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/p2p/net/mock/mock_notif_test.go b/p2p/net/mock/mock_notif_test.go
index 488e86d..3521d6f 100644
--- a/p2p/net/mock/mock_notif_test.go
+++ b/p2p/net/mock/mock_notif_test.go
@@ -11,8 +11,9 @@ import (
 )
 
 func TestNotifications(t *testing.T) {
+	const swarmSize = 5
 
-	mn, err := FullMeshLinked(context.Background(), 5)
+	mn, err := FullMeshLinked(context.Background(), swarmSize)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -23,7 +24,7 @@ func TestNotifications(t *testing.T) {
 	nets := mn.Nets()
 	notifiees := make([]*netNotifiee, len(nets))
 	for i, pn := range nets {
-		n := newNetNotifiee()
+		n := newNetNotifiee(swarmSize)
 		pn.Notify(n)
 		notifiees[i] = n
 	}
@@ -193,14 +194,14 @@ type netNotifiee struct {
 	closedStream chan inet.Stream
 }
 
-func newNetNotifiee() *netNotifiee {
+func newNetNotifiee(buffer int) *netNotifiee {
 	return &netNotifiee{
-		listen:       make(chan ma.Multiaddr),
-		listenClose:  make(chan ma.Multiaddr),
-		connected:    make(chan inet.Conn),
-		disconnected: make(chan inet.Conn),
-		openedStream: make(chan inet.Stream),
-		closedStream: make(chan inet.Stream),
+		listen:       make(chan ma.Multiaddr, buffer),
+		listenClose:  make(chan ma.Multiaddr, buffer),
+		connected:    make(chan inet.Conn, buffer),
+		disconnected: make(chan inet.Conn, buffer),
+		openedStream: make(chan inet.Stream, buffer),
+		closedStream: make(chan inet.Stream, buffer),
 	}
 }
 
diff --git a/p2p/net/mock/mock_peernet.go b/p2p/net/mock/mock_peernet.go
index d0265b8..26d5a3b 100644
--- a/p2p/net/mock/mock_peernet.go
+++ b/p2p/net/mock/mock_peernet.go
@@ -31,7 +31,7 @@ type peernet struct {
 	streamHandler inet.StreamHandler
 	connHandler   inet.ConnHandler
 
-	notifmu sync.RWMutex
+	notifmu sync.Mutex
 	notifs  map[inet.Notifiee]struct{}
 
 	proc goprocess.Process
@@ -381,11 +381,17 @@ func (pn *peernet) StopNotify(f inet.Notifiee) {
 
 // notifyAll runs the notification function on all Notifiees
 func (pn *peernet) notifyAll(notification func(f inet.Notifiee)) {
-	pn.notifmu.RLock()
+	pn.notifmu.Lock()
+	var wg sync.WaitGroup
 	for n := range pn.notifs {
 		// make sure we dont block
 		// and they dont block each other.
-		go notification(n)
+		wg.Add(1)
+		go func(n inet.Notifiee) {
+			defer wg.Done()
+			notification(n)
+		}(n)
 	}
-	pn.notifmu.RUnlock()
+	wg.Wait()
+	pn.notifmu.Unlock()
 }
-- 
GitLab