diff --git a/floodsub_test.go b/floodsub_test.go
index c15c57d..47df0f5 100644
--- a/floodsub_test.go
+++ b/floodsub_test.go
@@ -1075,6 +1075,8 @@ func TestSubscriptionNotification(t *testing.T) {
 
 	msgs := make([]*Subscription, numHosts)
 	subPeersFound := make([]map[peer.ID]struct{}, numHosts)
+
+	wg := sync.WaitGroup{}
 	for i, ps := range psubs {
 		subch, err := ps.Subscribe("foobar")
 		if err != nil {
@@ -1084,7 +1086,9 @@ func TestSubscriptionNotification(t *testing.T) {
 		msgs[i] = subch
 		peersFound := make(map[peer.ID]struct{})
 		subPeersFound[i] = peersFound
+		wg.Add(1)
 		go func(peersFound map[peer.ID]struct{}) {
+			defer wg.Done()
 			for i := 0; i < numHosts-1; i++ {
 				pid, err := subch.NextSubscribedPeer(ctx)
 				if err != nil {
@@ -1099,6 +1103,7 @@ func TestSubscriptionNotification(t *testing.T) {
 
 	time.Sleep(time.Millisecond * 100)
 
+	wg.Wait()
 	for _, peersFound := range subPeersFound {
 		if len(peersFound) != numHosts-1 {
 			t.Fatal("incorrect number of peers found")