From ee5a3a826e217072752d2fead4bfdf5656683853 Mon Sep 17 00:00:00 2001
From: Roman Tkachenko <roman@gravitational.com>
Date: Mon, 19 Jul 2021 10:22:19 +0300
Subject: [PATCH] Fix test

---
 lib/srv/app/server_test.go | 54 ++++++++++++++++++++++++--------------
 1 file changed, 35 insertions(+), 19 deletions(-)

diff --git a/lib/srv/app/server_test.go b/lib/srv/app/server_test.go
index 5be39ada10a62..add05d47f3224 100644
--- a/lib/srv/app/server_test.go
+++ b/lib/srv/app/server_test.go
@@ -120,6 +120,21 @@ func (s *Suite) TearDownSuite(c *check.C) {
 	c.Assert(err, check.IsNil)
 }
 
+var (
+	staticLabels = map[string]string{
+		"bar": "baz",
+	}
+	dynamicLabelName    = "qux"
+	dynamicLabelPeriod  = types.NewDuration(time.Second)
+	dynamicLabelCommand = []string{"expr", "1", "+", "3"}
+	dynamicLabels       = map[string]types.CommandLabel{
+		dynamicLabelName: &types.CommandLabelV2{
+			Period:  dynamicLabelPeriod,
+			Command: dynamicLabelCommand,
+		},
+	}
+)
+
 func (s *Suite) SetUpTest(c *check.C) {
 	s.closeContext, s.closeFunc = context.WithCancel(context.Background())
 
@@ -140,15 +155,6 @@ func (s *Suite) SetUpTest(c *check.C) {
 	s.hostport = u.Host
 
 	// Create a services.App that will be used for each test.
-	staticLabels := map[string]string{
-		"bar": "baz",
-	}
-	dynamicLabels := map[string]types.CommandLabel{
-		"qux": &types.CommandLabelV2{
-			Period:  types.NewDuration(time.Second),
-			Command: []string{"expr", "1", "+", "3"},
-		},
-	}
 	s.hostUUID = uuid.New()
 	s.server = &types.ServerV2{
 		Kind:    types.KindAppServer,
@@ -270,18 +276,28 @@ func (s *Suite) TestStart(c *check.C) {
 
 	// Check that the services.Server sent via heartbeat is correct. For example,
 	// check that the dynamic labels have been evaluated.
-	c.Assert(server.GetApps(), check.HasLen, 1)
-	app := server.GetApps()[0]
 
-	c.Assert(app.Name, check.Equals, "foo")
-	c.Assert(app.URI, check.Equals, s.testhttp.URL)
-	c.Assert(app.PublicAddr, check.Equals, "foo.example.com")
-	c.Assert(app.StaticLabels, check.DeepEquals, map[string]string{
-		"bar": "baz",
+	c.Assert(server.GetApps(), check.DeepEquals, []*types.App{
+		{
+			Name:         "foo",
+			URI:          s.testhttp.URL,
+			PublicAddr:   "foo.example.com",
+			StaticLabels: staticLabels,
+			DynamicLabels: map[string]types.CommandLabelV2{
+				dynamicLabelName: {
+					Period:  dynamicLabelPeriod,
+					Command: dynamicLabelCommand,
+					Result:  "4",
+				},
+			},
+		},
+		{
+			Name:         "awsconsole",
+			URI:          constants.AWSConsoleURL,
+			PublicAddr:   "aws.example.com",
+			StaticLabels: staticLabels,
+		},
 	})
-	dynamicLabel, ok := app.DynamicLabels["qux"]
-	c.Assert(ok, check.Equals, true)
-	c.Assert(dynamicLabel.GetResult(), check.Equals, "4")
 
 	// Check the expiry time is correct.
 	c.Assert(s.clock.Now().Before(server.Expiry()), check.Equals, true)