refactor: introduce a helper to reduce duplication

weaveworks · Dec 8, 2015 · 1b6d156 · 1b6d156
1 parent 88f9f0a
commit 1b6d156
Showing 1 changed file with 13 additions and 18 deletions.
diff --git a/mesh/gossip_test.go b/mesh/gossip_test.go
@@ -115,12 +115,19 @@ func checkTopology(t *testing.T, router *Router, wantedPeers ...*Peer) {
 	router.Peers.RUnlock()
 }
 
+func flushAndCheckTopology(t *testing.T, routers []*Router, wantedPeers ...*Peer) {
+	sendPendingGossip(routers...)
+	for _, r := range routers {
+		checkTopology(t, r, wantedPeers...)
+	}
+}
+
 func implTestGossipTopology(t *testing.T) {
 	// Create some peers that will talk to each other
 	r1 := NewTestRouter("01:00:00:01:00:00")
 	r2 := NewTestRouter("02:00:00:02:00:00")
 	r3 := NewTestRouter("03:00:00:03:00:00")
-
+	routers := []*Router{r1, r2, r3}
 	// Check state when they have no connections
 	checkTopology(t, r1, r1.tp())
 	checkTopology(t, r2, r2.tp())
@@ -132,23 +139,14 @@ func implTestGossipTopology(t *testing.T) {
 	checkTopology(t, r2, r1.tp(r2), r2.tp(r1))
 
 	AddTestGossipConnection(r2, r3)
-	sendPendingGossip(r1, r2, r3)
-	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
-	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
-	checkTopology(t, r3, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
+	flushAndCheckTopology(t, routers, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
 
 	AddTestGossipConnection(r3, r1)
-	sendPendingGossip(r1, r2, r3)
-	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1, r2))
-	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1, r2))
-	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1, r2))
+	flushAndCheckTopology(t, routers, r1.tp(r2, r3), r2.tp(r1, r3), r3.tp(r1, r2))
 
 	// Drop the connection from 2 to 3
 	r2.DeleteTestGossipConnection(r3)
-	sendPendingGossip(r1, r2, r3)
-	checkTopology(t, r1, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1, r2))
-	checkTopology(t, r2, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1, r2))
-	checkTopology(t, r3, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1, r2))
+	flushAndCheckTopology(t, routers, r1.tp(r2, r3), r2.tp(r1), r3.tp(r1, r2))
 
 	// Drop the connection from 1 to 3
 	r1.DeleteTestGossipConnection(r3)
@@ -160,7 +158,6 @@ func implTestGossipTopology(t *testing.T) {
 
 	// On a timer, r3 will gossip to r1
 	r3.SendAllGossip()
-	sendPendingGossip(r1, r2, r3)
 	checkTopology(t, r1, r1.tp(r2), r2.tp(r1), r3.tp(r1, r2))
 }
 
@@ -169,12 +166,10 @@ func TestGossipSurrogate(t *testing.T) {
 	r1 := NewTestRouter("01:00:00:01:00:00")
 	r2 := NewTestRouter("02:00:00:02:00:00")
 	r3 := NewTestRouter("03:00:00:03:00:00")
+	routers := []*Router{r1, r2, r3}
 	AddTestGossipConnection(r1, r2)
 	AddTestGossipConnection(r3, r2)
-	sendPendingGossip(r1, r2, r3)
-	checkTopology(t, r1, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
-	checkTopology(t, r2, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
-	checkTopology(t, r3, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
+	flushAndCheckTopology(t, routers, r1.tp(r2), r2.tp(r1, r3), r3.tp(r2))
 
 	// create a gossiper at either end, but not the middle
 	g1 := newTestGossiper()