From 75314cb910a7edf7e1ab90f029eb21bc1e581391 Mon Sep 17 00:00:00 2001
From: Mike Lang <mikelang3000@gmail.com>
Date: Mon, 3 Apr 2017 10:42:33 -0700
Subject: [PATCH] Reduce manually listing all topologies in a few places

Prefer WalkTopologies to apply a uniform action to every topology,
reducing need to make multiple changes and risk of errors if you forget one.
---
 probe/appclient/app_client_internal_test.go | 28 +++------------------
 probe/probe_internal_test.go                | 15 +++--------
 probe/topology_tagger.go                    | 13 +---------
 3 files changed, 8 insertions(+), 48 deletions(-)

diff --git a/probe/appclient/app_client_internal_test.go b/probe/appclient/app_client_internal_test.go
index 73edd4074..66bcae51d 100644
--- a/probe/appclient/app_client_internal_test.go
+++ b/probe/appclient/app_client_internal_test.go
@@ -73,30 +73,10 @@ func TestAppClientPublish(t *testing.T) {
 	// marshalling->unmarshaling is not idempotent due to `json:"omitempty"`
 	// tags, transforming empty slices into nils. So, we make DeepEqual
 	// happy by setting empty `json:"omitempty"` entries to nil
-	rpt.Endpoint = report.MakeTopology()
-	rpt.Process = report.MakeTopology()
-	rpt.Container = report.MakeTopology()
-	rpt.ContainerImage = report.MakeTopology()
-	rpt.Pod = report.MakeTopology()
-	rpt.Service = report.MakeTopology()
-	rpt.Deployment = report.MakeTopology()
-	rpt.ReplicaSet = report.MakeTopology()
-	rpt.Host = report.MakeTopology()
-	rpt.Overlay = report.MakeTopology()
-	rpt.ECSTask = report.MakeTopology()
-	rpt.ECSService = report.MakeTopology()
-	rpt.Endpoint.Controls = nil
-	rpt.Process.Controls = nil
-	rpt.Container.Controls = nil
-	rpt.ContainerImage.Controls = nil
-	rpt.Pod.Controls = nil
-	rpt.Service.Controls = nil
-	rpt.Deployment.Controls = nil
-	rpt.ReplicaSet.Controls = nil
-	rpt.Host.Controls = nil
-	rpt.Overlay.Controls = nil
-	rpt.ECSTask.Controls = nil
-	rpt.ECSService.Controls = nil
+	rpt.WalkTopologies(func(to *report.Topology) {
+		*to = report.MakeTopology()
+		to.Controls = nil
+	})
 
 	s := dummyServer(t, token, id, version, rpt, done)
 	defer s.Close()
diff --git a/probe/probe_internal_test.go b/probe/probe_internal_test.go
index 6b6b4f005..42b7cbd58 100644
--- a/probe/probe_internal_test.go
+++ b/probe/probe_internal_test.go
@@ -81,18 +81,9 @@ func TestProbe(t *testing.T) {
 	// tags, transforming empty slices into nils. So, we make DeepEqual
 	// happy by setting empty `json:"omitempty"` entries to nil
 	node.Metrics = nil
-	want.Endpoint.Controls = nil
-	want.Process.Controls = nil
-	want.Container.Controls = nil
-	want.ContainerImage.Controls = nil
-	want.Pod.Controls = nil
-	want.Service.Controls = nil
-	want.Deployment.Controls = nil
-	want.ReplicaSet.Controls = nil
-	want.Host.Controls = nil
-	want.Overlay.Controls = nil
-	want.ECSTask.Controls = nil
-	want.ECSService.Controls = nil
+	want.WalkTopologies(func(t *report.Topology) {
+		t.Controls = nil
+	})
 	want.Endpoint.AddNode(node)
 
 	pub := mockPublisher{make(chan report.Report, 10)}
diff --git a/probe/topology_tagger.go b/probe/topology_tagger.go
index cf41896d1..ce37e19a3 100644
--- a/probe/topology_tagger.go
+++ b/probe/topology_tagger.go
@@ -16,18 +16,7 @@ func (topologyTagger) Name() string { return "Topology" }
 
 // Tag implements Tagger
 func (topologyTagger) Tag(r report.Report) (report.Report, error) {
-	for name, t := range map[string]*report.Topology{
-		report.Endpoint:       &(r.Endpoint),
-		report.Process:        &(r.Process),
-		report.Container:      &(r.Container),
-		report.ContainerImage: &(r.ContainerImage),
-		report.Pod:            &(r.Pod),
-		report.Service:        &(r.Service),
-		report.ECSTask:        &(r.ECSTask),
-		report.ECSService:     &(r.ECSService),
-		report.Host:           &(r.Host),
-		report.Overlay:        &(r.Overlay),
-	} {
+	for name, t := range r.TopologyMap() {
 		for _, node := range t.Nodes {
 			t.AddNode(node.WithTopology(name))
 		}