From 273a0fba474a784137c20e7658b478b26a107953 Mon Sep 17 00:00:00 2001
From: Bryan Boreham <bjboreham@gmail.com>
Date: Wed, 15 Jul 2015 12:33:11 +0100
Subject: [PATCH] Add sender to mock broadcast messages

---
 ipam/allocator_test.go  |  4 ++--
 ipam/testutils_test.go  |  2 +-
 testing/gossip/mocks.go | 12 ++++++------
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/ipam/allocator_test.go b/ipam/allocator_test.go
index 67d162b8..930dc437 100644
--- a/ipam/allocator_test.go
+++ b/ipam/allocator_test.go
@@ -254,9 +254,9 @@ func TestTransfer(t *testing.T) {
 	_, err = alloc3.Allocate("bar", subnet, nil)
 	require.True(t, err == nil, "Failed to get address")
 
-	router.GossipBroadcast(alloc2.Gossip())
+	alloc2.gossip.GossipBroadcast(alloc2.Gossip())
 	router.Flush()
-	router.GossipBroadcast(alloc3.Gossip())
+	alloc2.gossip.GossipBroadcast(alloc3.Gossip())
 	router.Flush()
 	router.RemovePeer(alloc2.ourName)
 	router.RemovePeer(alloc3.ourName)
diff --git a/ipam/testutils_test.go b/ipam/testutils_test.go
index 294444d6..95f6407a 100644
--- a/ipam/testutils_test.go
+++ b/ipam/testutils_test.go
@@ -208,7 +208,7 @@ func makeNetworkOfAllocators(size int, cidr string) ([]*Allocator, *gossip.TestR
 		allocs[i] = alloc
 	}
 
-	gossipRouter.GossipBroadcast(allocs[size-1].Gossip())
+	allocs[size-1].gossip.GossipBroadcast(allocs[size-1].Gossip())
 	gossipRouter.Flush()
 	return allocs, gossipRouter, subnet
 }
diff --git a/testing/gossip/mocks.go b/testing/gossip/mocks.go
index 64889d29..43f7ee5a 100644
--- a/testing/gossip/mocks.go
+++ b/testing/gossip/mocks.go
@@ -34,10 +34,10 @@ func NewTestRouter(loss float32) *TestRouter {
 	return &TestRouter{make(map[router.PeerName]chan interface{}, 100), loss}
 }
 
-func (grouter *TestRouter) GossipBroadcast(update router.GossipData) error {
+func (grouter *TestRouter) gossipBroadcast(sender router.PeerName, update router.GossipData) error {
 	for _, gossipChan := range grouter.gossipChans {
 		select {
-		case gossipChan <- broadcastMessage{data: update}:
+		case gossipChan <- broadcastMessage{sender: sender, data: update}:
 		default: // drop the message if we cannot send it
 			common.Log.Errorf("Dropping message")
 		}
@@ -66,7 +66,7 @@ type TestRouterClient struct {
 	sender router.PeerName
 }
 
-func (grouter *TestRouter) run(gossiper router.Gossiper, gossipChan chan interface{}) {
+func (grouter *TestRouter) run(sender router.PeerName, gossiper router.Gossiper, gossipChan chan interface{}) {
 	gossipTimer := time.Tick(10 * time.Second)
 	for {
 		select {
@@ -98,7 +98,7 @@ func (grouter *TestRouter) run(gossiper router.Gossiper, gossipChan chan interfa
 				}
 			}
 		case <-gossipTimer:
-			grouter.GossipBroadcast(gossiper.Gossip())
+			grouter.gossipBroadcast(sender, gossiper.Gossip())
 		}
 	}
 }
@@ -106,7 +106,7 @@ func (grouter *TestRouter) run(gossiper router.Gossiper, gossipChan chan interfa
 func (grouter *TestRouter) Connect(sender router.PeerName, gossiper router.Gossiper) router.Gossip {
 	gossipChan := make(chan interface{}, 100)
 
-	go grouter.run(gossiper, gossipChan)
+	go grouter.run(sender, gossiper, gossipChan)
 
 	grouter.gossipChans[sender] = gossipChan
 	return TestRouterClient{grouter, sender}
@@ -122,5 +122,5 @@ func (client TestRouterClient) GossipUnicast(dstPeerName router.PeerName, buf []
 }
 
 func (client TestRouterClient) GossipBroadcast(update router.GossipData) error {
-	return client.router.GossipBroadcast(update)
+	return client.router.gossipBroadcast(client.sender, update)
 }