discovery: aws/ec2 unit tests (#14364)

* discovery: add aws/ec2 unit tests * discovery: initial skeleton for aws/ec2 unit tests This is a - very likely - not too useful unit test for the AWS SD. It is commited so other people can check the basic logic and the implementation. Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: fix linter complains about ec2_test.go Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: add basic unit test for aws This tests only the basic labelling, not including the VPC related information. Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: fix linter complains about ec2_test.go Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: other linter fixes in aws/ec2_test.go Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: implement remaining tests for aws/ec2 The coverage is not 100% but I think it is a good starting point if someone wants to improve that. Currently it covers all the AWS API calls. Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: make linter happy in aws/ec2_test.go Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: make utility funtcions private Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discover: no global variable in the aws/ec2 test Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: common body for some tests in ec2 Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: try to make golangci-lint happy Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: make every non-test function private Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: test for errors first in TestRefresh Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: move refresh tests into the function This way people can find both the test cases and the execution of the test at the same place. Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: fix copyright date Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: remove misleading comment Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: rename test for easier identification Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: use static values for the test cases Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discover: try to make the linter happy Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: drop redundant data from ec2 and use common ptr functions Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: use Error instead of Equal Signed-off-by: Arpad Kunszt <akunszt@hiya.com> * discovery: merge refreshAZIDs tests into one Signed-off-by: Arpad Kunszt <akunszt@hiya.com> --------- Signed-off-by: Arpad Kunszt <akunszt@hiya.com>
1 month ago · 08a7162502
2 changed files with 437 additions and 2 deletions
--- a/discovery/aws/ec2.go
+++ b/discovery/aws/ec2.go
@ -30,6 +30,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws/ec2metadata"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/config"
 	"github.com/prometheus/common/model"
@ -148,7 +149,7 @@ type EC2Discovery struct {
 	*refresh.Discovery
 	logger *slog.Logger
 	cfg    *EC2SDConfig
-	ec2    *ec2.EC2
+	ec2    ec2iface.EC2API

 	// azToAZID maps this account's availability zones to their underlying AZ
 	// ID, e.g. eu-west-2a -> euw2-az2. Refreshes are performed sequentially, so
@ -182,7 +183,7 @@ func NewEC2Discovery(conf *EC2SDConfig, logger *slog.Logger, metrics discovery.D
 	return d, nil
 }

-func (d *EC2Discovery) ec2Client(context.Context) (*ec2.EC2, error) {
+func (d *EC2Discovery) ec2Client(context.Context) (ec2iface.EC2API, error) {
 	if d.ec2 != nil {
 		return d.ec2, nil
 	}
--- a/discovery/aws/ec2_test.go
+++ b/discovery/aws/ec2_test.go
@ -0,0 +1,434 @@
+// Copyright 2024 The Prometheus Authors
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package aws
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
+	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
+
+	"github.com/prometheus/prometheus/discovery/targetgroup"
+)
+
+// Helper function to get pointers on literals.
+// NOTE: this is common between a few tests. In the future it might worth to move this out into a separate package.
+func strptr(str string) *string {
+	return &str
+}
+
+func boolptr(b bool) *bool {
+	return &b
+}
+
+func int64ptr(i int64) *int64 {
+	return &i
+}
+
+// Struct for test data.
+type ec2DataStore struct {
+	region string
+
+	azToAZID map[string]string
+
+	ownerID string
+
+	instances []*ec2.Instance
+}
+
+// The tests itself.
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
+
+func TestEC2DiscoveryRefreshAZIDs(t *testing.T) {
+	ctx := context.Background()
+
+	// iterate through the test cases
+	for _, tt := range []struct {
+		name       string
+		shouldFail bool
+		ec2Data    *ec2DataStore
+	}{
+		{
+			name:       "Normal",
+			shouldFail: false,
+			ec2Data: &ec2DataStore{
+				azToAZID: map[string]string{
+					"azname-a": "azid-1",
+					"azname-b": "azid-2",
+					"azname-c": "azid-3",
+				},
+			},
+		},
+		{
+			name:       "HandleError",
+			shouldFail: true,
+			ec2Data:    &ec2DataStore{},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			client := newMockEC2Client(tt.ec2Data)
+
+			d := &EC2Discovery{
+				ec2: client,
+			}
+
+			err := d.refreshAZIDs(ctx)
+			if tt.shouldFail {
+				require.Error(t, err)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, client.ec2Data.azToAZID, d.azToAZID)
+			}
+		})
+	}
+}
+
+func TestEC2DiscoveryRefresh(t *testing.T) {
+	ctx := context.Background()
+
+	// iterate through the test cases
+	for _, tt := range []struct {
+		name     string
+		ec2Data  *ec2DataStore
+		expected []*targetgroup.Group
+	}{
+		{
+			name: "NoPrivateIp",
+			ec2Data: &ec2DataStore{
+				region: "region-noprivateip",
+				azToAZID: map[string]string{
+					"azname-a": "azid-1",
+					"azname-b": "azid-2",
+					"azname-c": "azid-3",
+				},
+				instances: []*ec2.Instance{
+					{
+						InstanceId: strptr("instance-id-noprivateip"),
+					},
+				},
+			},
+			expected: []*targetgroup.Group{
+				{
+					Source: "region-noprivateip",
+				},
+			},
+		},
+		{
+			name: "NoVpc",
+			ec2Data: &ec2DataStore{
+				region: "region-novpc",
+				azToAZID: map[string]string{
+					"azname-a": "azid-1",
+					"azname-b": "azid-2",
+					"azname-c": "azid-3",
+				},
+				ownerID: "owner-id-novpc",
+				instances: []*ec2.Instance{
+					{
+						// set every possible options and test them here
+						Architecture:      strptr("architecture-novpc"),
+						ImageId:           strptr("ami-novpc"),
+						InstanceId:        strptr("instance-id-novpc"),
+						InstanceLifecycle: strptr("instance-lifecycle-novpc"),
+						InstanceType:      strptr("instance-type-novpc"),
+						Placement:         &ec2.Placement{AvailabilityZone: strptr("azname-b")},
+						Platform:          strptr("platform-novpc"),
+						PrivateDnsName:    strptr("private-dns-novpc"),
+						PrivateIpAddress:  strptr("1.2.3.4"),
+						PublicDnsName:     strptr("public-dns-novpc"),
+						PublicIpAddress:   strptr("42.42.42.2"),
+						State:             &ec2.InstanceState{Name: strptr("running")},
+						// test tags once and for all
+						Tags: []*ec2.Tag{
+							{Key: strptr("tag-1-key"), Value: strptr("tag-1-value")},
+							{Key: strptr("tag-2-key"), Value: strptr("tag-2-value")},
+							nil,
+							{Value: strptr("tag-4-value")},
+							{Key: strptr("tag-5-key")},
+						},
+					},
+				},
+			},
+			expected: []*targetgroup.Group{
+				{
+					Source: "region-novpc",
+					Targets: []model.LabelSet{
+						{
+							"__address__":                     model.LabelValue("1.2.3.4:4242"),
+							"__meta_ec2_ami":                  model.LabelValue("ami-novpc"),
+							"__meta_ec2_architecture":         model.LabelValue("architecture-novpc"),
+							"__meta_ec2_availability_zone":    model.LabelValue("azname-b"),
+							"__meta_ec2_availability_zone_id": model.LabelValue("azid-2"),
+							"__meta_ec2_instance_id":          model.LabelValue("instance-id-novpc"),
+							"__meta_ec2_instance_lifecycle":   model.LabelValue("instance-lifecycle-novpc"),
+							"__meta_ec2_instance_type":        model.LabelValue("instance-type-novpc"),
+							"__meta_ec2_instance_state":       model.LabelValue("running"),
+							"__meta_ec2_owner_id":             model.LabelValue("owner-id-novpc"),
+							"__meta_ec2_platform":             model.LabelValue("platform-novpc"),
+							"__meta_ec2_private_dns_name":     model.LabelValue("private-dns-novpc"),
+							"__meta_ec2_private_ip":           model.LabelValue("1.2.3.4"),
+							"__meta_ec2_public_dns_name":      model.LabelValue("public-dns-novpc"),
+							"__meta_ec2_public_ip":            model.LabelValue("42.42.42.2"),
+							"__meta_ec2_region":               model.LabelValue("region-novpc"),
+							"__meta_ec2_tag_tag_1_key":        model.LabelValue("tag-1-value"),
+							"__meta_ec2_tag_tag_2_key":        model.LabelValue("tag-2-value"),
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "Ipv4",
+			ec2Data: &ec2DataStore{
+				region: "region-ipv4",
+				azToAZID: map[string]string{
+					"azname-a": "azid-1",
+					"azname-b": "azid-2",
+					"azname-c": "azid-3",
+				},
+				instances: []*ec2.Instance{
+					{
+						// just the minimum needed for the refresh work
+						ImageId:          strptr("ami-ipv4"),
+						InstanceId:       strptr("instance-id-ipv4"),
+						InstanceType:     strptr("instance-type-ipv4"),
+						Placement:        &ec2.Placement{AvailabilityZone: strptr("azname-c")},
+						PrivateIpAddress: strptr("5.6.7.8"),
+						State:            &ec2.InstanceState{Name: strptr("running")},
+						SubnetId:         strptr("azid-3"),
+						VpcId:            strptr("vpc-ipv4"),
+						// network intefaces
+						NetworkInterfaces: []*ec2.InstanceNetworkInterface{
+							// interface without subnet -> should be ignored
+							{
+								Ipv6Addresses: []*ec2.InstanceIpv6Address{
+									{
+										Ipv6Address:   strptr("2001:db8:1::1"),
+										IsPrimaryIpv6: boolptr(true),
+									},
+								},
+							},
+							// interface with subnet, no IPv6
+							{
+								Ipv6Addresses: []*ec2.InstanceIpv6Address{},
+								SubnetId:      strptr("azid-3"),
+							},
+							// interface with another subnet, no IPv6
+							{
+								Ipv6Addresses: []*ec2.InstanceIpv6Address{},
+								SubnetId:      strptr("azid-1"),
+							},
+						},
+					},
+				},
+			},
+			expected: []*targetgroup.Group{
+				{
+					Source: "region-ipv4",
+					Targets: []model.LabelSet{
+						{
+							"__address__":                     model.LabelValue("5.6.7.8:4242"),
+							"__meta_ec2_ami":                  model.LabelValue("ami-ipv4"),
+							"__meta_ec2_availability_zone":    model.LabelValue("azname-c"),
+							"__meta_ec2_availability_zone_id": model.LabelValue("azid-3"),
+							"__meta_ec2_instance_id":          model.LabelValue("instance-id-ipv4"),
+							"__meta_ec2_instance_state":       model.LabelValue("running"),
+							"__meta_ec2_instance_type":        model.LabelValue("instance-type-ipv4"),
+							"__meta_ec2_owner_id":             model.LabelValue(""),
+							"__meta_ec2_primary_subnet_id":    model.LabelValue("azid-3"),
+							"__meta_ec2_private_ip":           model.LabelValue("5.6.7.8"),
+							"__meta_ec2_region":               model.LabelValue("region-ipv4"),
+							"__meta_ec2_subnet_id":            model.LabelValue(",azid-3,azid-1,"),
+							"__meta_ec2_vpc_id":               model.LabelValue("vpc-ipv4"),
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "Ipv6",
+			ec2Data: &ec2DataStore{
+				region: "region-ipv6",
+				azToAZID: map[string]string{
+					"azname-a": "azid-1",
+					"azname-b": "azid-2",
+					"azname-c": "azid-3",
+				},
+				instances: []*ec2.Instance{
+					{
+						// just the minimum needed for the refresh work
+						ImageId:          strptr("ami-ipv6"),
+						InstanceId:       strptr("instance-id-ipv6"),
+						InstanceType:     strptr("instance-type-ipv6"),
+						Placement:        &ec2.Placement{AvailabilityZone: strptr("azname-b")},
+						PrivateIpAddress: strptr("9.10.11.12"),
+						State:            &ec2.InstanceState{Name: strptr("running")},
+						SubnetId:         strptr("azid-2"),
+						VpcId:            strptr("vpc-ipv6"),
+						// network intefaces
+						NetworkInterfaces: []*ec2.InstanceNetworkInterface{
+							// interface without primary IPv6, index 2
+							{
+								Attachment: &ec2.InstanceNetworkInterfaceAttachment{
+									DeviceIndex: int64ptr(3),
+								},
+								Ipv6Addresses: []*ec2.InstanceIpv6Address{
+									{
+										Ipv6Address:   strptr("2001:db8:2::1:1"),
+										IsPrimaryIpv6: boolptr(false),
+									},
+								},
+								SubnetId: strptr("azid-2"),
+							},
+							// interface with primary IPv6, index 1
+							{
+								Attachment: &ec2.InstanceNetworkInterfaceAttachment{
+									DeviceIndex: int64ptr(1),
+								},
+								Ipv6Addresses: []*ec2.InstanceIpv6Address{
+									{
+										Ipv6Address:   strptr("2001:db8:2::2:1"),
+										IsPrimaryIpv6: boolptr(false),
+									},
+									{
+										Ipv6Address:   strptr("2001:db8:2::2:2"),
+										IsPrimaryIpv6: boolptr(true),
+									},
+								},
+								SubnetId: strptr("azid-2"),
+							},
+							// interface with primary IPv6, index 3
+							{
+								Attachment: &ec2.InstanceNetworkInterfaceAttachment{
+									DeviceIndex: int64ptr(3),
+								},
+								Ipv6Addresses: []*ec2.InstanceIpv6Address{
+									{
+										Ipv6Address:   strptr("2001:db8:2::3:1"),
+										IsPrimaryIpv6: boolptr(true),
+									},
+								},
+								SubnetId: strptr("azid-1"),
+							},
+							// interface without primary IPv6, index 0
+							{
+								Attachment: &ec2.InstanceNetworkInterfaceAttachment{
+									DeviceIndex: int64ptr(0),
+								},
+								Ipv6Addresses: []*ec2.InstanceIpv6Address{},
+								SubnetId:      strptr("azid-3"),
+							},
+						},
+					},
+				},
+			},
+			expected: []*targetgroup.Group{
+				{
+					Source: "region-ipv6",
+					Targets: []model.LabelSet{
+						{
+							"__address__":                       model.LabelValue("9.10.11.12:4242"),
+							"__meta_ec2_ami":                    model.LabelValue("ami-ipv6"),
+							"__meta_ec2_availability_zone":      model.LabelValue("azname-b"),
+							"__meta_ec2_availability_zone_id":   model.LabelValue("azid-2"),
+							"__meta_ec2_instance_id":            model.LabelValue("instance-id-ipv6"),
+							"__meta_ec2_instance_state":         model.LabelValue("running"),
+							"__meta_ec2_instance_type":          model.LabelValue("instance-type-ipv6"),
+							"__meta_ec2_ipv6_addresses":         model.LabelValue(",2001:db8:2::1:1,2001:db8:2::2:1,2001:db8:2::2:2,2001:db8:2::3:1,"),
+							"__meta_ec2_owner_id":               model.LabelValue(""),
+							"__meta_ec2_primary_ipv6_addresses": model.LabelValue(",,2001:db8:2::2:2,,2001:db8:2::3:1,"),
+							"__meta_ec2_primary_subnet_id":      model.LabelValue("azid-2"),
+							"__meta_ec2_private_ip":             model.LabelValue("9.10.11.12"),
+							"__meta_ec2_region":                 model.LabelValue("region-ipv6"),
+							"__meta_ec2_subnet_id":              model.LabelValue(",azid-2,azid-1,azid-3,"),
+							"__meta_ec2_vpc_id":                 model.LabelValue("vpc-ipv6"),
+						},
+					},
+				},
+			},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			client := newMockEC2Client(tt.ec2Data)
+
+			d := &EC2Discovery{
+				ec2: client,
+				cfg: &EC2SDConfig{
+					Port:   4242,
+					Region: client.ec2Data.region,
+				},
+			}
+
+			g, err := d.refresh(ctx)
+			require.NoError(t, err)
+			require.Equal(t, tt.expected, g)
+		})
+	}
+}
+
+// EC2 client mock.
+type mockEC2Client struct {
+	ec2iface.EC2API
+	ec2Data ec2DataStore
+}
+
+func newMockEC2Client(ec2Data *ec2DataStore) *mockEC2Client {
+	client := mockEC2Client{
+		ec2Data: *ec2Data,
+	}
+	return &client
+}
+
+func (m *mockEC2Client) DescribeAvailabilityZonesWithContext(ctx aws.Context, input *ec2.DescribeAvailabilityZonesInput, opts ...request.Option) (*ec2.DescribeAvailabilityZonesOutput, error) {
+	if len(m.ec2Data.azToAZID) == 0 {
+		return nil, errors.New("No AZs found")
+	}
+
+	azs := make([]*ec2.AvailabilityZone, len(m.ec2Data.azToAZID))
+
+	i := 0
+	for k, v := range m.ec2Data.azToAZID {
+		azs[i] = &ec2.AvailabilityZone{
+			ZoneName: strptr(k),
+			ZoneId:   strptr(v),
+		}
+		i++
+	}
+
+	return &ec2.DescribeAvailabilityZonesOutput{
+		AvailabilityZones: azs,
+	}, nil
+}
+
+func (m *mockEC2Client) DescribeInstancesPagesWithContext(ctx aws.Context, input *ec2.DescribeInstancesInput, fn func(*ec2.DescribeInstancesOutput, bool) bool, opts ...request.Option) error {
+	r := ec2.Reservation{}
+	r.SetInstances(m.ec2Data.instances)
+	r.SetOwnerId(m.ec2Data.ownerID)
+
+	o := ec2.DescribeInstancesOutput{}
+	o.SetReservations([]*ec2.Reservation{&r})
+
+	_ = fn(&o, true)
+
+	return nil
+}