From 7cec30f0e100b8f4515fe0c37a8761f9978c1c1e Mon Sep 17 00:00:00 2001
From: Shirley <4163034+fridgepoet@users.noreply.github.com>
Date: Wed, 23 Nov 2022 17:13:47 +0100
Subject: [PATCH] CloudWatch: Fix flaky test assertions (#59234)

---
 .../cloudwatch/routes/dimension_keys_test.go    | 15 ++++++++-------
 .../cloudwatch/routes/dimension_values_test.go  | 17 +++++++++--------
 2 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/pkg/tsdb/cloudwatch/routes/dimension_keys_test.go b/pkg/tsdb/cloudwatch/routes/dimension_keys_test.go
index e48fa7dabb6..068e0fd4bca 100644
--- a/pkg/tsdb/cloudwatch/routes/dimension_keys_test.go
+++ b/pkg/tsdb/cloudwatch/routes/dimension_keys_test.go
@@ -24,7 +24,14 @@ var logger = &logtest.Fake{}
 func Test_DimensionKeys_Route(t *testing.T) {
 	t.Run("calls FilterDimensionKeysRequest when a StandardDimensionKeysRequest is passed", func(t *testing.T) {
 		mockListMetricsService := mocks.ListMetricsServiceMock{}
-		mockListMetricsService.On("GetDimensionKeysByDimensionFilter", mock.Anything).Return([]string{}, nil).Once()
+		mockListMetricsService.On("GetDimensionKeysByDimensionFilter", mock.MatchedBy(func(r resources.DimensionKeysRequest) bool {
+			return r.ResourceRequest != nil && *r.ResourceRequest == resources.ResourceRequest{Region: "us-east-2"} &&
+				r.Namespace == "AWS/EC2" &&
+				r.MetricName == "CPUUtilization" &&
+				len(r.DimensionFilter) == 2 &&
+				assert.Contains(t, r.DimensionFilter, &resources.Dimension{Name: "NodeID", Value: "Shared"}) &&
+				assert.Contains(t, r.DimensionFilter, &resources.Dimension{Name: "stage", Value: "QueryCommit"})
+		})).Return([]string{}, nil).Once()
 		newListMetricsService = func(pluginCtx backend.PluginContext, reqCtxFactory models.RequestContextFactoryFunc, region string) (models.ListMetricsProvider, error) {
 			return &mockListMetricsService, nil
 		}
@@ -32,12 +39,6 @@ func Test_DimensionKeys_Route(t *testing.T) {
 		req := httptest.NewRequest("GET", `/dimension-keys?region=us-east-2&namespace=AWS/EC2&metricName=CPUUtilization&dimensionFilters={"NodeID":["Shared"],"stage":["QueryCommit"]}`, nil)
 		handler := http.HandlerFunc(ResourceRequestMiddleware(DimensionKeysHandler, logger, nil))
 		handler.ServeHTTP(rr, req)
-		mockListMetricsService.AssertCalled(t, "GetDimensionKeysByDimensionFilter", resources.DimensionKeysRequest{
-			ResourceRequest: &resources.ResourceRequest{Region: "us-east-2"},
-			Namespace:       "AWS/EC2",
-			MetricName:      "CPUUtilization",
-			DimensionFilter: []*resources.Dimension{{Name: "NodeID", Value: "Shared"}, {Name: "stage", Value: "QueryCommit"}},
-		})
 	})
 
 	t.Run("calls GetHardCodedDimensionKeysByNamespace when a StandardDimensionKeysRequest is passed", func(t *testing.T) {
diff --git a/pkg/tsdb/cloudwatch/routes/dimension_values_test.go b/pkg/tsdb/cloudwatch/routes/dimension_values_test.go
index c2853a8c98b..cf365c1a090 100644
--- a/pkg/tsdb/cloudwatch/routes/dimension_values_test.go
+++ b/pkg/tsdb/cloudwatch/routes/dimension_values_test.go
@@ -18,7 +18,15 @@ import (
 func Test_DimensionValues_Route(t *testing.T) {
 	t.Run("Calls GetDimensionValuesByDimensionFilter when a valid request is passed", func(t *testing.T) {
 		mockListMetricsService := mocks.ListMetricsServiceMock{}
-		mockListMetricsService.On("GetDimensionValuesByDimensionFilter", mock.Anything).Return([]string{}, nil).Once()
+		mockListMetricsService.On("GetDimensionValuesByDimensionFilter", mock.MatchedBy(func(r resources.DimensionValuesRequest) bool {
+			return r.ResourceRequest != nil && *r.ResourceRequest == resources.ResourceRequest{Region: "us-east-2"} &&
+				r.Namespace == "AWS/EC2" &&
+				r.MetricName == "CPUUtilization" &&
+				r.DimensionKey == "instanceId" &&
+				len(r.DimensionFilter) == 2 &&
+				assert.Contains(t, r.DimensionFilter, &resources.Dimension{Name: "NodeID", Value: "Shared"}) &&
+				assert.Contains(t, r.DimensionFilter, &resources.Dimension{Name: "stage", Value: "QueryCommit"})
+		})).Return([]string{}, nil).Once()
 		newListMetricsService = func(pluginCtx backend.PluginContext, reqCtxFactory models.RequestContextFactoryFunc, region string) (models.ListMetricsProvider, error) {
 			return &mockListMetricsService, nil
 		}
@@ -26,13 +34,6 @@ func Test_DimensionValues_Route(t *testing.T) {
 		req := httptest.NewRequest("GET", `/dimension-values?region=us-east-2&dimensionKey=instanceId&namespace=AWS/EC2&metricName=CPUUtilization&dimensionFilters={"NodeID":["Shared"],"stage":["QueryCommit"]}`, nil)
 		handler := http.HandlerFunc(ResourceRequestMiddleware(DimensionValuesHandler, logger, nil))
 		handler.ServeHTTP(rr, req)
-		mockListMetricsService.AssertCalled(t, "GetDimensionValuesByDimensionFilter", resources.DimensionValuesRequest{
-			ResourceRequest: &resources.ResourceRequest{Region: "us-east-2"},
-			Namespace:       "AWS/EC2",
-			MetricName:      "CPUUtilization",
-			DimensionKey:    "instanceId",
-			DimensionFilter: []*resources.Dimension{{Name: "NodeID", Value: "Shared"}, {Name: "stage", Value: "QueryCommit"}},
-		})
 	})
 
 	t.Run("returns 500 if GetDimensionValuesByDimensionFilter returns an error", func(t *testing.T) {