Add sender to mock broadcast messages

dvusboy · Jul 15, 2015 · 273a0fb · 273a0fb
1 parent acf13dc
commit 273a0fb
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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,15 +98,15 @@ func (grouter *TestRouter) run(gossiper router.Gossiper, gossipChan chan interfa
 				}
 			}
 		case <-gossipTimer:
-			grouter.GossipBroadcast(gossiper.Gossip())
+			grouter.gossipBroadcast(sender, gossiper.Gossip())
 		}
 	}
 }
 
 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)
 }