From f796615e10348520111edf1566b17b3b812192d3 Mon Sep 17 00:00:00 2001
From: Juan Batiz-Benet <juan@benet.ai>
Date: Fri, 26 Jun 2015 23:57:01 -0700
Subject: [PATCH] mock: fix notif test

License: MIT
Signed-off-by: Juan Batiz-Benet <juan@benet.ai>
---
 p2p/net/mock/mock_notif_test.go | 36 +++++++++++++++++++--------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/p2p/net/mock/mock_notif_test.go b/p2p/net/mock/mock_notif_test.go
index 1ef9fc6aa..d91403f8b 100644
--- a/p2p/net/mock/mock_notif_test.go
+++ b/p2p/net/mock/mock_notif_test.go
@@ -43,24 +43,30 @@ func TestNotifications(t *testing.T) {
 	for i, s := range nets {
 		n := notifiees[i]
 		for _, s2 := range nets {
-			cos := s.ConnsToPeer(s2.LocalPeer())
-			func() {
-				for i := 0; i < len(cos); i++ {
-					var c inet.Conn
-					select {
-					case c = <-n.connected:
-					case <-time.After(timeout):
-						t.Fatal("timeout")
-					}
-					for _, c2 := range cos {
-						if c == c2 {
-							t.Log("got notif for conn")
-							return
-						}
+			var actual []inet.Conn
+			for len(s.ConnsToPeer(s2.LocalPeer())) != len(actual) {
+				select {
+				case c := <-n.connected:
+					actual = append(actual, c)
+				case <-time.After(timeout):
+					t.Fatal("timeout")
+				}
+			}
+
+			expect := s.ConnsToPeer(s2.LocalPeer())
+			for _, c1 := range actual {
+				found := false
+				for _, c2 := range expect {
+					if c1 == c2 {
+						found = true
+						break
 					}
+				}
+				if !found {
 					t.Error("connection not found")
 				}
-			}()
+			}
+
 		}
 	}