From c659e3448ad551996ff2a235a43f20efea3a2da4 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Thu, 25 Jun 2015 17:38:02 +0100
Subject: [PATCH] Add a flush() function to the mock gossip router, and call it
 where required

---
 ipam/allocator_test.go |  4 ++++
 ipam/testutils_test.go | 16 +++++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/ipam/allocator_test.go b/ipam/allocator_test.go
index 2aaece78..6cc9aae5 100644
--- a/ipam/allocator_test.go
+++ b/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))
 
diff --git a/ipam/testutils_test.go b/ipam/testutils_test.go
index fb2af1fd..8448c2a2 100644
--- a/ipam/testutils_test.go
+++ b/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
 }