From 1c1a45a8809dcb4bfaa3012a4a98f774fc8c2d0a Mon Sep 17 00:00:00 2001
From: Shirley <4163034+fridgepoet@users.noreply.github.com>
Date: Tue, 15 Feb 2022 12:45:50 +0100
Subject: [PATCH] CloudWatch: Add tests to annotations query (#45337)

---
 pkg/tsdb/cloudwatch/annotation_query_test.go | 106 +++++++++++++++++++
 pkg/tsdb/cloudwatch/test_utils.go            |  25 +++++
 2 files changed, 131 insertions(+)
 create mode 100644 pkg/tsdb/cloudwatch/annotation_query_test.go

diff --git a/pkg/tsdb/cloudwatch/annotation_query_test.go b/pkg/tsdb/cloudwatch/annotation_query_test.go
new file mode 100644
index 00000000000..6f3d7c71d6b
--- /dev/null
+++ b/pkg/tsdb/cloudwatch/annotation_query_test.go
@@ -0,0 +1,106 @@
+package cloudwatch
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/session"
+	"github.com/aws/aws-sdk-go/service/cloudwatch"
+	"github.com/aws/aws-sdk-go/service/cloudwatch/cloudwatchiface"
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestQuery_AnnotationQuery(t *testing.T) {
+	origNewCWClient := NewCWClient
+	t.Cleanup(func() {
+		NewCWClient = origNewCWClient
+	})
+
+	var client FakeCWAnnotationsClient
+	NewCWClient = func(sess *session.Session) cloudwatchiface.CloudWatchAPI {
+		return &client
+	}
+
+	t.Run("DescribeAlarmsForMetric is called with minimum parameters", func(t *testing.T) {
+		client = FakeCWAnnotationsClient{describeAlarmsForMetricOutput: &cloudwatch.DescribeAlarmsForMetricOutput{}}
+		im := datasource.NewInstanceManager(func(s backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
+			return datasourceInfo{}, nil
+		})
+
+		executor := newExecutor(im, newTestConfig(), fakeSessionCache{})
+		_, err := executor.QueryData(context.Background(), &backend.QueryDataRequest{
+			PluginContext: backend.PluginContext{
+				DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
+			},
+			Queries: []backend.DataQuery{
+				{
+					JSON: json.RawMessage(`{
+						"type":    "annotationQuery",
+						"region":    "us-east-1",
+						"namespace": "custom",
+						"metricName": "CPUUtilization",
+						"statistic": "Average"
+					}`),
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		require.Len(t, client.calls.describeAlarmsForMetric, 1)
+		assert.Equal(t, &cloudwatch.DescribeAlarmsForMetricInput{
+			Namespace:  aws.String("custom"),
+			MetricName: aws.String("CPUUtilization"),
+			Statistic:  aws.String("Average"),
+			Period:     aws.Int64(300),
+		}, client.calls.describeAlarmsForMetric[0])
+	})
+
+	t.Run("DescribeAlarms is called when prefixMatching is true", func(t *testing.T) {
+		client = FakeCWAnnotationsClient{describeAlarmsOutput: &cloudwatch.DescribeAlarmsOutput{}}
+		im := datasource.NewInstanceManager(func(s backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
+			return datasourceInfo{}, nil
+		})
+
+		executor := newExecutor(im, newTestConfig(), fakeSessionCache{})
+		_, err := executor.QueryData(context.Background(), &backend.QueryDataRequest{
+			PluginContext: backend.PluginContext{
+				DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{},
+			},
+			Queries: []backend.DataQuery{
+				{
+					JSON: json.RawMessage(`{
+						"type":    "annotationQuery",
+						"region":    "us-east-1",
+						"namespace": "custom",
+						"metricName": "CPUUtilization",
+						"statistic": "Average",
+						"prefixMatching": true,
+						"actionPrefix": "some_action_prefix",
+						"alarmNamePrefix": "some_alarm_name_prefix"
+					}`),
+				},
+			},
+		})
+		require.NoError(t, err)
+
+		require.Len(t, client.calls.describeAlarms, 1)
+		assert.Equal(t, &cloudwatch.DescribeAlarmsInput{
+			MaxRecords:      pointerInt64(100),
+			ActionPrefix:    pointerString("some_action_prefix"),
+			AlarmNamePrefix: pointerString("some_alarm_name_prefix"),
+		}, client.calls.describeAlarms[0])
+	})
+}
+
+func pointerString(s string) *string {
+	return &s
+}
+func pointerInt64(i int64) *int64 {
+	return &i
+}
diff --git a/pkg/tsdb/cloudwatch/test_utils.go b/pkg/tsdb/cloudwatch/test_utils.go
index a6727f74ca0..8eac803c83a 100644
--- a/pkg/tsdb/cloudwatch/test_utils.go
+++ b/pkg/tsdb/cloudwatch/test_utils.go
@@ -79,6 +79,31 @@ func (c FakeCWClient) ListMetricsPages(input *cloudwatch.ListMetricsInput, fn fu
 	return nil
 }
 
+type FakeCWAnnotationsClient struct {
+	cloudwatchiface.CloudWatchAPI
+	calls annontationsQueryCalls
+
+	describeAlarmsForMetricOutput *cloudwatch.DescribeAlarmsForMetricOutput
+	describeAlarmsOutput          *cloudwatch.DescribeAlarmsOutput
+}
+
+type annontationsQueryCalls struct {
+	describeAlarmsForMetric []*cloudwatch.DescribeAlarmsForMetricInput
+	describeAlarms          []*cloudwatch.DescribeAlarmsInput
+}
+
+func (c *FakeCWAnnotationsClient) DescribeAlarmsForMetric(params *cloudwatch.DescribeAlarmsForMetricInput) (*cloudwatch.DescribeAlarmsForMetricOutput, error) {
+	c.calls.describeAlarmsForMetric = append(c.calls.describeAlarmsForMetric, params)
+
+	return c.describeAlarmsForMetricOutput, nil
+}
+
+func (c *FakeCWAnnotationsClient) DescribeAlarms(params *cloudwatch.DescribeAlarmsInput) (*cloudwatch.DescribeAlarmsOutput, error) {
+	c.calls.describeAlarms = append(c.calls.describeAlarms, params)
+
+	return c.describeAlarmsOutput, nil
+}
+
 type fakeEC2Client struct {
 	ec2iface.EC2API