Add tests for the ec2 fetcher

lib/cloud/watchers/watcher_test.go

@@ -28,6 +28,7 @@ import (
 	"github.com/gravitational/teleport/lib/srv/db/cloud"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/elasticache"
 	"github.com/aws/aws-sdk-go/service/memorydb"
 	"github.com/aws/aws-sdk-go/service/rds"
@@ -93,11 +94,27 @@ func TestWatcher(t *testing.T) {
 		aws.StringValue(memorydbUnsupported.ARN): memorydbUnsupportedTags,
 	}
 
+	ec2InstanceDiscovered := makeEC2Instances(
+		"instance-1",
+		ec2.InstanceStateNameRunning,
+		map[string]string{"teleport": "yes"})
+
+	ec2InstanceIgnored := makeEC2Instances(
+		"instance-1",
+		ec2.InstanceStateNameRunning,
+		map[string]string{"teleport": "no"})
+
+	ec2InstanceOff := makeEC2Instances(
+		"instance-1",
+		ec2.InstanceStateNamePending,
+		map[string]string{"teleport": "yes"})
+
 	tests := []struct {
-		name              string
-		awsMatchers       []services.AWSMatcher
-		clients           clients.CloudClients
-		expectedDatabases types.Databases
+		name                 string
+		awsMatchers          []services.AWSMatcher
+		clients              clients.CloudClients
+		expectedDatabases    types.Databases
+		expectedEC2Instances []EC2Instances
 	}{
 		{
 			name: "RDS labels matching",
@@ -294,6 +311,30 @@ func TestWatcher(t *testing.T) {
 				memorydbDatabaseProd,
 			},
 		},
+		{
+			name: "EC2",
+			awsMatchers: []services.AWSMatcher{
+				{
+					Types:   []string{services.AWSMatcherEC2},
+					Regions: []string{"us-east-1"},
+					Tags:    types.Labels{"teleport": []string{"yes"}},
+				},
+			},
+			clients: &clients.TestCloudClients{
+				EC2: &cloud.EC2Mock{
+					Instances: []*ec2.Instance{
+						ec2InstanceDiscovered, // labels match
+						ec2InstanceIgnored,    // labels do not match
+						ec2InstanceOff,
+					},
+				},
+			},
+			expectedEC2Instances: []EC2Instances{
+				EC2Instances{
+					Instances: []*ec2.Instance{ec2InstanceDiscovered},
+				},
+			},
+		},
 	}
 
 	for _, test := range tests {
@@ -307,6 +348,8 @@ func TestWatcher(t *testing.T) {
 				// makeFetchers function uses a map for matcher types so
 				// databases can come in random orders.
 				require.ElementsMatch(t, test.expectedDatabases, databases)
+			case ec2Instances := <-watcher.EC2C():
+				require.ElementsMatch(t, test.expectedEC2Instances, ec2Instances)
 			case <-time.After(time.Second):
 				t.Fatal("didn't receive databases after 1 second")
 			}
@@ -525,3 +568,21 @@ func labelsToTags(labels map[string]string) (tags []*rds.Tag) {
 	}
 	return tags
 }
+
+func makeEC2Instances(instID, state string, labels map[string]string) *ec2.Instance {
+	var tags []*ec2.Tag
+	for key, val := range labels {
+		tags = append(tags, &ec2.Tag{
+			Key:   aws.String(key),
+			Value: aws.String(val),
+		})
+	}
+
+	return &ec2.Instance{
+		InstanceId: &instID,
+		State: &ec2.InstanceState{
+			Name: aws.String(state),
+		},
+		Tags: tags,
+	}
+}

lib/srv/db/cloud/mocks.go

@@ -19,10 +19,13 @@ package cloud
 import (
 	"context"
 	"crypto/tls"
+	"strings"
 	"sync"
 
 	"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/aws/aws-sdk-go/service/elasticache"
 	"github.com/aws/aws-sdk-go/service/elasticache/elasticacheiface"
 	"github.com/aws/aws-sdk-go/service/iam"
@@ -512,3 +515,44 @@ func (m *MemoryDBMock) UpdateUserWithContext(_ aws.Context, input *memorydb.Upda
 	}
 	return nil, trace.NotFound("user %s not found", aws.StringValue(input.UserName))
 }
+
+type EC2Mock struct {
+	ec2iface.EC2API
+	Instances []*ec2.Instance
+}
+
+func (m *EC2Mock) DescribeInstancesPagesWithContext(
+	ctx context.Context, input *ec2.DescribeInstancesInput,
+	f func(dio *ec2.DescribeInstancesOutput, b bool) bool, opts ...request.Option) error {
+
+	var instances []*ec2.Instance
+
+	for _, inst := range m.Instances {
+		tagMatch := false
+		stateMatch := false
+		for _, tag := range inst.Tags {
+			for _, filter := range input.Filters {
+				if strings.HasPrefix(aws.StringValue(filter.Name), "tag:") && !tagMatch {
+					tagMatch =
+						aws.StringValue(filter.Name)[4:] == aws.StringValue(tag.Key) &&
+							aws.StringValue(tag.Value) == aws.StringValueSlice(filter.Values)[0]
+				}
+				if aws.StringValue(filter.Name) == "instance-state-name" && !stateMatch {
+					stateMatch =
+						aws.StringValue(inst.State.Name) == ec2.InstanceStateNameRunning
+				}
+
+				if stateMatch && tagMatch {
+					instances = append(instances, inst)
+				}
+			}
+		}
+
+	}
+
+	filtered := &ec2.DescribeInstancesOutput{
+		Reservations: []*ec2.Reservation{{Instances: instances}},
+	}
+	f(filtered, true)
+	return nil
+}