From bf9c27a166372027b58f4c5c51fba86709e90793 Mon Sep 17 00:00:00 2001
From: Janos Guljas <janos@resenje.org>
Date: Fri, 20 Jul 2018 14:46:53 +0200
Subject: [PATCH] swarm/network: wait for subs in PeerEvents and fix
 stream.runSyncTest

---
 swarm/network/simulation/events.go         | 11 +++++++++++
 swarm/network/stream/snapshot_sync_test.go | 19 ++++++++-----------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/swarm/network/simulation/events.go b/swarm/network/simulation/events.go
index f9cfadb73e..4db8424908 100644
--- a/swarm/network/simulation/events.go
+++ b/swarm/network/simulation/events.go
@@ -18,6 +18,7 @@ package simulation
 
 import (
 	"context"
+	"sync"
 
 	"github.com/ethereum/go-ethereum/p2p/discover"
 
@@ -71,24 +72,32 @@ func (f *PeerEventsFilter) MsgCode(c uint64) *PeerEventsFilter {
 func (s *Simulation) PeerEvents(ctx context.Context, ids []discover.NodeID, filters ...*PeerEventsFilter) <-chan PeerEvent {
 	eventC := make(chan PeerEvent)
 
+	// wait group to make sure all subscriptons to admin peerEvents are established
+	// before this function returns.
+	var subsWG sync.WaitGroup
 	for _, id := range ids {
 		s.shutdownWG.Add(1)
+		subsWG.Add(1)
 		go func(id discover.NodeID) {
 			defer s.shutdownWG.Done()
 
 			client, err := s.Net.GetNode(id).Client()
 			if err != nil {
+				subsWG.Done()
 				eventC <- PeerEvent{NodeID: id, Error: err}
 				return
 			}
 			events := make(chan *p2p.PeerEvent)
 			sub, err := client.Subscribe(ctx, "admin", events, "peerEvents")
 			if err != nil {
+				subsWG.Done()
 				eventC <- PeerEvent{NodeID: id, Error: err}
 				return
 			}
 			defer sub.Unsubscribe()
 
+			subsWG.Done()
+
 			for {
 				select {
 				case <-ctx.Done():
@@ -153,5 +162,7 @@ func (s *Simulation) PeerEvents(ctx context.Context, ids []discover.NodeID, filt
 		}(id)
 	}
 
+	// wait all subscriptions
+	subsWG.Wait()
 	return eventC
 }
diff --git a/swarm/network/stream/snapshot_sync_test.go b/swarm/network/stream/snapshot_sync_test.go
index 63dd6c305e..c9aacee02c 100644
--- a/swarm/network/stream/snapshot_sync_test.go
+++ b/swarm/network/stream/snapshot_sync_test.go
@@ -384,18 +384,15 @@ func runSyncTest(chunkCount int, nodeCount int) error {
 			subscriptionCount += cnt
 		}
 
-		go func() {
-			for e := range eventC {
-				if e.Error != nil {
-					//return e.Error
-					return
-				}
-				subscriptionCount--
-				if subscriptionCount == 0 {
-					break
-				}
+		for e := range eventC {
+			if e.Error != nil {
+				return e.Error
 			}
-		}()
+			subscriptionCount--
+			if subscriptionCount == 0 {
+				break
+			}
+		}
 		//get the the node at that index
 		//this is the node selected for upload
 		node := sim.RandomUpNode()