From 134c3c0ed9825395676bd5ba8cadf18e2d0e853b Mon Sep 17 00:00:00 2001
From: johncming <johncming@yahoo.com>
Date: Fri, 20 Dec 2019 22:27:58 +0800
Subject: [PATCH] move walkRoute to dispatch package. (#2136)

Signed-off-by: johncming <johncming@yahoo.com>
---
 cmd/alertmanager/main.go | 15 +------
 dispatch/route.go        | 11 ++++++
 dispatch/route_test.go   | 85 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 98 insertions(+), 13 deletions(-)

diff --git a/cmd/alertmanager/main.go b/cmd/alertmanager/main.go
index 0469c393..aefaf59e 100644
--- a/cmd/alertmanager/main.go
+++ b/cmd/alertmanager/main.go
@@ -172,17 +172,6 @@ func buildReceiverIntegrations(nc *config.Receiver, tmpl *template.Template, log
 	return integrations, nil
 }
 
-// walkRoute traverses the route tree in depth-first order.
-func walkRoute(r *dispatch.Route, visit func(*dispatch.Route)) {
-	visit(r)
-	if r.Routes == nil {
-		return
-	}
-	for i := range r.Routes {
-		walkRoute(r.Routes[i], visit)
-	}
-}
-
 func main() {
 	os.Exit(run())
 }
@@ -403,7 +392,7 @@ func run() int {
 		// Build the routing tree and record which receivers are used.
 		routes := dispatch.NewRoute(conf.Route, nil)
 		activeReceivers := make(map[string]struct{})
-		walkRoute(routes, func(r *dispatch.Route) {
+		routes.Walk(func(r *dispatch.Route) {
 			activeReceivers[r.RouteOpts.Receiver] = struct{}{}
 		})
 
@@ -447,7 +436,7 @@ func run() int {
 		})
 
 		disp = dispatch.NewDispatcher(alerts, routes, pipeline, marker, timeoutFunc, logger, dispMetrics)
-		walkRoute(routes, func(r *dispatch.Route) {
+		routes.Walk(func(r *dispatch.Route) {
 			if r.RouteOpts.RepeatInterval > *retention {
 				level.Warn(configLogger).Log(
 					"msg",
diff --git a/dispatch/route.go b/dispatch/route.go
index 43c48c25..b4606b7f 100644
--- a/dispatch/route.go
+++ b/dispatch/route.go
@@ -155,6 +155,17 @@ func (r *Route) Key() string {
 	return b.String()
 }
 
+// Walk traverses the route tree in depth-first order.
+func (r *Route) Walk(visit func(*Route)) {
+	visit(r)
+	if r.Routes == nil {
+		return
+	}
+	for i := range r.Routes {
+		r.Routes[i].Walk(visit)
+	}
+}
+
 // RouteOpts holds various routing options necessary for processing alerts
 // that match a given route.
 type RouteOpts struct {
diff --git a/dispatch/route_test.go b/dispatch/route_test.go
index e291b3f2..fdfa7f3d 100644
--- a/dispatch/route_test.go
+++ b/dispatch/route_test.go
@@ -267,3 +267,88 @@ routes:
 		}
 	}
 }
+
+func TestRouteWalk(t *testing.T) {
+	in := `
+receiver: 'notify-def'
+
+routes:
+- match:
+    owner: 'team-A'
+
+  receiver: 'notify-A'
+
+  routes:
+  - match:
+      env: 'testing'
+
+    receiver: 'notify-testing'
+    group_by: [...]
+
+  - match:
+      env: "production"
+
+    receiver: 'notify-productionA'
+    group_wait: 1m
+
+    continue: true
+
+  - match_re:
+      env: "produ.*"
+      job: ".*"
+
+    receiver: 'notify-productionB'
+    group_wait: 30s
+    group_interval: 5m
+    repeat_interval: 1h
+    group_by: ['job']
+
+
+- match_re:
+    owner: 'team-(B|C)'
+
+  group_by: ['foo', 'bar']
+  group_wait: 2m
+  receiver: 'notify-BC'
+
+- match:
+    group_by: 'role'
+  group_by: ['role']
+
+  routes:
+  - match:
+      env: 'testing'
+    receiver: 'notify-testing'
+    routes:
+    - match:
+        wait: 'long'
+      group_wait: 2m
+`
+
+	var ctree config.Route
+	if err := yaml.UnmarshalStrict([]byte(in), &ctree); err != nil {
+		t.Fatal(err)
+	}
+	tree := NewRoute(&ctree, nil)
+
+	expected := []string{
+		"notify-def",
+		"notify-A",
+		"notify-testing",
+		"notify-productionA",
+		"notify-productionB",
+		"notify-BC",
+		"notify-def",
+		"notify-testing",
+		"notify-testing",
+	}
+
+	var got []string
+	tree.Walk(func(r *Route) {
+		got = append(got, r.RouteOpts.Receiver)
+	})
+
+	if !reflect.DeepEqual(got, expected) {
+		t.Errorf("\nexpected:\n%v\ngot:\n%v", expected, got)
+	}
+}