fix same test in swarm

License: MIT
Signed-off-by: Jeromy <jeromyj@gmail.com>

p2p/net/mock/mock_notif_test.go

@@ -43,28 +43,42 @@ func TestNotifications(t *testing.T) {
 	// test everyone got the correct connection opened calls
 	for i, s := range nets {
 		n := notifiees[i]
-		notifs := make(map[peer.ID]inet.Conn)
-		for j := 0; j < len(nets)-1; j++ {
-			select {
-			case c := <-n.connected:
-				_, ok := notifs[c.RemotePeer()]
-				if ok {
-					t.Fatal("shouldnt have received more than one connection per peer")
+		notifs := make(map[peer.ID][]inet.Conn)
+		for j, s2 := range nets {
+			if i == j {
+				continue
+			}
+
+			// this feels a little sketchy, but its probably okay
+			for len(s.ConnsToPeer(s2.LocalPeer())) != len(notifs[s2.LocalPeer()]) {
+				select {
+				case c := <-n.connected:
+					nfp := notifs[c.RemotePeer()]
+					notifs[c.RemotePeer()] = append(nfp, c)
+				case <-time.After(timeout):
+					t.Fatal("timeout")
 				}
-				notifs[c.RemotePeer()] = c
-			case <-time.After(timeout):
-				t.Fatal("timeout")
 			}
 		}
 
-		for p, con := range notifs {
+		for p, cons := range notifs {
 			expect := s.ConnsToPeer(p)
-			if len(expect) != 1 {
-				t.Fatal("got more than one connection, not supposed to happen")
+			if len(expect) != len(cons) {
+				t.Fatal("got different number of connections")
 			}
 
-			if expect[0] != con {
-				t.Fatal("got different connection than we expected")
+			for _, c := range cons {
+				var found bool
+				for _, c2 := range expect {
+					if c == c2 {
+						found = true
+						break
+					}
+				}
+
+				if !found {
+					t.Fatal("connection not found!")
+				}
 			}
 		}
 	}

p2p/net/swarm/swarm_notif_test.go

@@ -8,6 +8,7 @@ import (
 	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
 
 	inet "github.com/ipfs/go-ipfs/p2p/net"
+	peer "github.com/ipfs/go-ipfs/p2p/peer"
 )
 
 func TestNotifications(t *testing.T) {
@@ -37,35 +38,43 @@ func TestNotifications(t *testing.T) {
 	// test everyone got the correct connection opened calls
 	for i, s := range swarms {
 		n := notifiees[i]
-		for _, s2 := range swarms {
-			if s == s2 {
+		notifs := make(map[peer.ID][]inet.Conn)
+		for j, s2 := range swarms {
+			if i == j {
 				continue
 			}
 
-			var actual []inet.Conn
-			for len(s.ConnectionsToPeer(s2.LocalPeer())) != len(actual) {
+			// this feels a little sketchy, but its probably okay
+			for len(s.ConnectionsToPeer(s2.LocalPeer())) != len(notifs[s2.LocalPeer()]) {
 				select {
 				case c := <-n.connected:
-					actual = append(actual, c)
+					nfp := notifs[c.RemotePeer()]
+					notifs[c.RemotePeer()] = append(nfp, c)
 				case <-time.After(timeout):
 					t.Fatal("timeout")
 				}
 			}
+		}
+
+		for p, cons := range notifs {
+			expect := s.ConnectionsToPeer(p)
+			if len(expect) != len(cons) {
+				t.Fatal("got different number of connections")
+			}
 
-			expect := s.ConnectionsToPeer(s2.LocalPeer())
-			for _, c1 := range actual {
-				found := false
+			for _, c := range cons {
+				var found bool
 				for _, c2 := range expect {
-					if c1 == c2 {
+					if c == c2 {
 						found = true
 						break
 					}
 				}
+
 				if !found {
-					t.Error("connection not found")
+					t.Fatal("connection not found!")
 				}
 			}
-
 		}
 	}