From 5535f5e738772a1ec310efc2daaaa45c0cc8a0ce Mon Sep 17 00:00:00 2001
From: Paul Bellamy <paul.a.bellamy@gmail.com>
Date: Fri, 19 Feb 2016 11:07:00 +0000
Subject: [PATCH] Review Feedback

---
 probe/docker/container.go |  6 +++++-
 render/topologies.go      | 19 ++++++-------------
 render/topologies_test.go | 20 ++++++++++++++++++++
 3 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/probe/docker/container.go b/probe/docker/container.go
index c2f02f4a7d..b25813a4d2 100644
--- a/probe/docker/container.go
+++ b/probe/docker/container.go
@@ -355,10 +355,14 @@ func (c *container) GetNode(hostID string, localAddrs []net.IP) report.Node {
 		result = result.WithControls(UnpauseContainer)
 	} else if c.container.State.Running {
 		uptime := (mtime.Now().Sub(c.container.State.StartedAt) / time.Second) * time.Second
+		networkMode := ""
+		if c.container.HostConfig != nil {
+			networkMode = c.container.HostConfig.NetworkMode
+		}
 		result = result.WithLatests(map[string]string{
 			ContainerUptime:       uptime.String(),
 			ContainerRestartCount: strconv.Itoa(c.container.RestartCount),
-			ContainerNetworkMode:  c.container.HostConfig.NetworkMode,
+			ContainerNetworkMode:  networkMode,
 		})
 		result = result.WithControls(
 			RestartContainer, StopContainer, PauseContainer, AttachContainer, ExecContainer,
diff --git a/render/topologies.go b/render/topologies.go
index 1adb663fa0..496787e371 100644
--- a/render/topologies.go
+++ b/render/topologies.go
@@ -137,35 +137,28 @@ func (r containerWithHostIPsRenderer) Render(rpt report.Report) RenderableNodes
 			continue
 		}
 
-		ips, ok := c.Node.Sets.Lookup(docker.ContainerIPs)
-		if ok && len(ips) > 0 {
-			continue
-		}
-
 		h, ok := hosts[MakeHostID(report.ExtractHostID(c.Node))]
 		if !ok {
 			continue
 		}
 
-		hostNetworks, ok := h.Sets.Lookup(host.LocalNetworks)
-		if !ok {
-			continue
-		}
-
-		hostIPs := report.MakeStringSet()
+		newIPs := report.MakeStringSet()
+		hostNetworks, _ := h.Sets.Lookup(host.LocalNetworks)
 		for _, cidr := range hostNetworks {
 			if ip, _, err := net.ParseCIDR(cidr); err == nil {
-				hostIPs = hostIPs.Add(ip.String())
+				newIPs = newIPs.Add(ip.String())
 			}
 		}
 
-		c.Sets = c.Sets.Add(docker.ContainerIPs, hostIPs)
+		c.Sets = c.Sets.Add(docker.ContainerIPs, newIPs)
 		containers[id] = c
 	}
 
 	return containers
 }
 
+// ContainerWithHostIPsRenderer is a Renderer which produces a container graph
+// enriched with host IPs on containers where NetworkMode is Host
 var ContainerWithHostIPsRenderer = containerWithHostIPsRenderer{ContainerRenderer}
 
 type containerWithImageNameRenderer struct {
diff --git a/render/topologies_test.go b/render/topologies_test.go
index 63ce03859e..bbbcfcdf2f 100644
--- a/render/topologies_test.go
+++ b/render/topologies_test.go
@@ -8,6 +8,7 @@ import (
 	"github.com/weaveworks/scope/probe/kubernetes"
 	"github.com/weaveworks/scope/render"
 	"github.com/weaveworks/scope/render/expected"
+	"github.com/weaveworks/scope/report"
 	"github.com/weaveworks/scope/test"
 	"github.com/weaveworks/scope/test/fixture"
 )
@@ -51,6 +52,25 @@ func TestContainerFilterRenderer(t *testing.T) {
 	}
 }
 
+func TestContainerWithHostIPsRenderer(t *testing.T) {
+	input := fixture.Report.Copy()
+	input.Container.Nodes[fixture.ClientContainerNodeID] = input.Container.Nodes[fixture.ClientContainerNodeID].WithLatests(map[string]string{
+		docker.ContainerNetworkMode: "host",
+	})
+	nodes := render.ContainerWithHostIPsRenderer.Render(input)
+
+	// Test host network nodes get the host IPs added.
+	haveNode, ok := nodes[render.MakeContainerID(fixture.ClientContainerID)]
+	if !ok {
+		t.Fatal("Expected output to have the client container node")
+	}
+	have, _ := haveNode.Sets.Lookup(docker.ContainerIPs)
+	want := report.MakeStringSet("10.10.10.0")
+	if !reflect.DeepEqual(want, have) {
+		t.Error(test.Diff(want, have))
+	}
+}
+
 func TestContainerFilterRendererImageName(t *testing.T) {
 	// Test nodes are filtered by image name as well.
 	input := fixture.Report.Copy()