Add a flush() function to the mock gossip router, and call it where required

ipam/allocator_test.go

@@ -256,13 +256,17 @@ func TestTransfer(t *testing.T) {
 	wt.AssertTrue(t, err == nil, "Failed to get address")
 
 	router.GossipBroadcast(alloc2.Gossip())
+	router.flush()
 	router.GossipBroadcast(alloc3.Gossip())
+	router.flush()
 	router.removePeer(alloc2.ourName)
 	router.removePeer(alloc3.ourName)
 	alloc2.Stop()
 	alloc3.Stop()
+	router.flush()
 	wt.AssertSuccess(t, alloc1.AdminTakeoverRanges(alloc2.ourName.String()))
 	wt.AssertSuccess(t, alloc1.AdminTakeoverRanges(alloc3.ourName.String()))
+	router.flush()
 
 	wt.AssertEquals(t, alloc1.NumFreeAddresses(subnet), address.Offset(1022))
 

ipam/testutils_test.go

@@ -210,6 +210,9 @@ type broadcastMessage struct {
 type exitMessage struct {
 	exitChan chan struct{}
 }
+type flushMessage struct {
+	flushChan chan struct{}
+}
 
 type TestGossipRouter struct {
 	gossipChans map[router.PeerName]chan interface{}
@@ -226,6 +229,14 @@ func (grouter *TestGossipRouter) GossipBroadcast(update router.GossipData) error
 	return nil
 }
 
+func (grouter *TestGossipRouter) flush() {
+	for _, gossipChan := range grouter.gossipChans {
+		flushChan := make(chan struct{})
+		gossipChan <- flushMessage{flushChan: flushChan}
+		<-flushChan
+	}
+}
+
 func (grouter *TestGossipRouter) removePeer(peer router.PeerName) {
 	gossipChan := grouter.gossipChans[peer]
 	resultChan := make(chan struct{})
@@ -251,6 +262,9 @@ func (grouter *TestGossipRouter) connect(sender router.PeerName, gossiper router
 					close(message.exitChan)
 					return
 				}
+				if message, isFlush := gossip.(flushMessage); isFlush {
+					close(message.flushChan)
+				}
 				if rand.Float32() > (1.0 - grouter.loss) {
 					continue
 				}
@@ -305,7 +319,7 @@ func makeNetworkOfAllocators(size int, cidr string) ([]*Allocator, TestGossipRou
 	}
 
 	gossipRouter.GossipBroadcast(allocs[size-1].Gossip())
-	time.Sleep(1000 * time.Millisecond)
+	gossipRouter.flush()
 	return allocs, gossipRouter, subnet
 }