Add logs agent integ test

test/cloudwatchlogs/publish_logs_test.go

@@ -6,6 +6,8 @@
 package cloudwatchlogs
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"log"
 	"os"
@@ -15,7 +17,13 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs"
 	"github.com/aws/aws-sdk-go-v2/service/cloudwatchlogs/types"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	ec2Types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/aws/amazon-cloudwatch-agent-test/environment"
@@ -29,9 +37,19 @@ const (
 	logLineId2                    = "bar"
 	logFilePath                   = "/tmp/cwagent_log_test.log" // TODO: not sure how well this will work on Windows
 	sleepForFlush                 = 20 * time.Second            // default flush interval is 5 seconds
+	sleepForExtendedFlush         = 180 * time.Second           // increase flush time for the two main tests
+	retryWaitTime                 = 30 * time.Second
 	configPathAutoRemoval         = "resources/config_auto_removal.json"
 	standardLogGroupClass         = "STANDARD"
 	infrequentAccessLogGroupClass = "INFREQUENT_ACCESS"
+	cwlPerfEndpoint               = "https://logs.us-west-2.amazonaws.com"
+	pdxRegionalCode               = "us-west-2"
+
+	entityType        = "@entity.KeyAttributes.Type"
+	entityName        = "@entity.KeyAttributes.Name"
+	entityEnvironment = "@entity.KeyAttributes.Environment"
+	entityPlatform    = "@entity.Attributes.PlatformType"
+	entityInstanceId  = "@entity.Attributes.EC2.InstanceId"
 )
 
 var (
@@ -70,6 +88,9 @@ var (
 			logGroupClass: types.LogGroupClassInfrequentAccess,
 		},
 	}
+	rnf       *types.ResourceNotFoundException
+	cwlClient *cloudwatchlogs.Client
+	ec2Client *ec2.Client
 )
 
 type writeToCloudWatchTestInput struct {
@@ -86,8 +107,29 @@ type cloudWatchLogGroupClassTestInput struct {
 	logGroupClass types.LogGroupClass
 }
 
+type expectedEntity struct {
+	entityType   string
+	name         string
+	environment  string
+	platformType string
+	instanceId   string
+}
+
 func init() {
 	environment.RegisterEnvironmentMetaDataFlags()
+	awsCfg, err := config.LoadDefaultConfig(
+		context.Background(),
+		config.WithRegion(pdxRegionalCode),
+	)
+	if err != nil {
+		fmt.Println("There was an error trying to load default config: ", err)
+		return
+	}
+
+	cwlClient = cloudwatchlogs.NewFromConfig(awsCfg, func(o *cloudwatchlogs.Options) {
+		o.BaseEndpoint = aws.String(cwlPerfEndpoint)
+	})
+	ec2Client = ec2.NewFromConfig(awsCfg)
 }
 
 // TestWriteLogsToCloudWatch writes N number of logs, and then validates that N logs
@@ -119,9 +161,9 @@ func TestWriteLogsToCloudWatch(t *testing.T) {
 
 			// ensure that there is enough time from the "start" time and the first log line,
 			// so we don't miss it in the GetLogEvents call
-			time.Sleep(sleepForFlush)
+			time.Sleep(sleepForExtendedFlush)
 			writeLogLines(t, f, param.iterations)
-			time.Sleep(sleepForFlush)
+			time.Sleep(sleepForExtendedFlush)
 			common.StopAgent()
 			end := time.Now()
 
@@ -139,6 +181,120 @@ func TestWriteLogsToCloudWatch(t *testing.T) {
 	}
 }
 
+// TestWriteLogsWithEntityInfo writes logs and validates that the
+// log events are associated with entities from CloudWatch Logs
+func TestWriteLogsWithEntityInfo(t *testing.T) {
+	instanceId := awsservice.GetInstanceId()
+	log.Printf("Found instance id %s", instanceId)
+
+	// Define tags to create for EC2 test case
+	tagsToCreate := []ec2Types.Tag{
+		{
+			Key:   aws.String("service"),
+			Value: aws.String("service-test"),
+		},
+	}
+
+	testCases := map[string]struct {
+		agentConfigPath string
+		iterations      int
+		useEC2Tag       bool
+		expectedEntity  expectedEntity
+	}{
+		"IAMRole": {
+			agentConfigPath: filepath.Join("resources", "config_log.json"),
+			iterations:      1000,
+			expectedEntity: expectedEntity{
+				entityType:   "Service",
+				name:         "cwa-e2e-iam-role", //should match the name of the IAM role used in our testing
+				environment:  "ec2:default",
+				platformType: "AWS::EC2",
+				instanceId:   instanceId,
+			},
+		},
+		"ServiceInConfig": {
+			agentConfigPath: filepath.Join("resources", "config_log_service_name.json"),
+			iterations:      1000,
+			expectedEntity: expectedEntity{
+				entityType:   "Service",
+				name:         "service-in-config",     //should match the service.name value in the config file
+				environment:  "environment-in-config", //should match the deployment.environment value in the config file
+				platformType: "AWS::EC2",
+				instanceId:   instanceId,
+			},
+		},
+		"EC2Tags": {
+			agentConfigPath: filepath.Join("resources", "config_log.json"),
+			iterations:      1000,
+			useEC2Tag:       true,
+			expectedEntity: expectedEntity{
+				entityType:   "Service",
+				name:         "service-test", //should match the value in tagsToCreate
+				environment:  "ec2:default",
+				platformType: "AWS::EC2",
+				instanceId:   instanceId,
+			},
+		},
+	}
+	for name, testCase := range testCases {
+		t.Run(name, func(t *testing.T) {
+			t.Cleanup(func() {
+				// delete the log group/stream after each test case
+				awsservice.DeleteLogGroupAndStream(instanceId, instanceId)
+
+				// delete EC2 tags added to the instance for the test
+				if testCase.useEC2Tag {
+					input := &ec2.DeleteTagsInput{
+						Resources: []string{instanceId},
+						Tags:      tagsToCreate,
+					}
+					_, err := ec2Client.DeleteTags(context.TODO(), input)
+					assert.NoError(t, err)
+					// Add a short delay to ensure tag deletion propagates
+					time.Sleep(5 * time.Second)
+				}
+			})
+			if testCase.useEC2Tag {
+				// enable instance metadata tags
+				modifyInput := &ec2.ModifyInstanceMetadataOptionsInput{
+					InstanceId:           aws.String(instanceId),
+					InstanceMetadataTags: ec2Types.InstanceMetadataTagsStateEnabled,
+				}
+				_, modifyErr := ec2Client.ModifyInstanceMetadataOptions(context.TODO(), modifyInput)
+				assert.NoError(t, modifyErr)
+
+				input := &ec2.CreateTagsInput{
+					Resources: []string{instanceId},
+					Tags:      tagsToCreate,
+				}
+				_, createErr := ec2Client.CreateTags(context.TODO(), input)
+				assert.NoError(t, createErr)
+			}
+			id := uuid.New()
+			f, err := os.Create(logFilePath + "-" + id.String())
+			if err != nil {
+				t.Fatalf("Error occurred creating log file for writing: %v", err)
+			}
+			common.DeleteFile(common.AgentLogFile)
+			common.TouchFile(common.AgentLogFile)
+
+			common.CopyFile(testCase.agentConfigPath, configOutputPath)
+
+			common.StartAgent(configOutputPath, true, false)
+			time.Sleep(sleepForExtendedFlush)
+			writeLogLines(t, f, testCase.iterations)
+			time.Sleep(sleepForExtendedFlush)
+			common.StopAgent()
+			end := time.Now()
+
+			ValidateEntity(t, instanceId, instanceId, &end, testCase.expectedEntity)
+
+			f.Close()
+			os.Remove(logFilePath + "-" + id.String())
+		})
+	}
+}
+
 // TestAutoRemovalStopAgent configures agent to monitor a file with auto removal on.
 // Then it restarts the agent.
 // Verify the file is NOT removed.
@@ -332,3 +488,152 @@ func checkData(t *testing.T, start time.Time, lineCount int) {
 	)
 	assert.NoError(t, err)
 }
+
+func ValidateEntity(t *testing.T, logGroup, logStream string, end *time.Time, expectedEntity expectedEntity) {
+	log.Printf("Validating entity for log group: %s, stream: %s", logGroup, logStream)
+
+	logGroupInfo, err := getLogGroup()
+	for _, lg := range logGroupInfo {
+		if *lg.LogGroupName == logGroup {
+			log.Println("Log group " + *lg.LogGroupName + " exists")
+			break
+		}
+	}
+	assert.NoError(t, err)
+	begin := end.Add(-sleepForExtendedFlush * 4)
+	log.Printf("Query start time is " + begin.String() + " and end time is " + end.String())
+	queryId, err := getLogQueryId(logGroup, &begin, end)
+	assert.NoError(t, err)
+	log.Printf("queryId is " + *queryId)
+	result, err := getQueryResult(queryId)
+	assert.NoError(t, err)
+	if !assert.NotZero(t, len(result)) {
+		return
+	}
+	requiredEntityFields := map[string]bool{
+		entityType:        false,
+		entityName:        false,
+		entityEnvironment: false,
+		entityPlatform:    false,
+		entityInstanceId:  false,
+	}
+	for _, field := range result[0] {
+		switch aws.ToString(field.Field) {
+		case entityType:
+			requiredEntityFields[entityType] = true
+			assert.Equal(t, expectedEntity.entityType, aws.ToString(field.Value))
+		case entityName:
+			requiredEntityFields[entityName] = true
+			assert.Equal(t, expectedEntity.name, aws.ToString(field.Value))
+		case entityEnvironment:
+			requiredEntityFields[entityEnvironment] = true
+			assert.Equal(t, expectedEntity.environment, aws.ToString(field.Value))
+		case entityPlatform:
+			requiredEntityFields[entityPlatform] = true
+			assert.Equal(t, expectedEntity.platformType, aws.ToString(field.Value))
+		case entityInstanceId:
+			requiredEntityFields[entityInstanceId] = true
+			assert.Equal(t, expectedEntity.instanceId, aws.ToString(field.Value))
+		}
+		fmt.Printf("%s: %s\n", aws.ToString(field.Field), aws.ToString(field.Value))
+	}
+	allEntityFieldsFound := true
+	for field, value := range requiredEntityFields {
+		if !value {
+			log.Printf("Missing required entity field: %s", field)
+			allEntityFieldsFound = false
+		}
+	}
+	assert.True(t, allEntityFieldsFound)
+}
+
+func getLogQueryId(logGroup string, since, until *time.Time) (*string, error) {
+	var queryId *string
+	params := &cloudwatchlogs.StartQueryInput{
+		QueryString:  aws.String("fields @message, @entity.KeyAttributes.Type, @entity.KeyAttributes.Name, @entity.KeyAttributes.Environment, @entity.Attributes.PlatformType, @entity.Attributes.EC2.InstanceId"),
+		LogGroupName: aws.String(logGroup),
+	}
+	if since != nil {
+		params.StartTime = aws.Int64(since.UnixMilli())
+	}
+	if until != nil {
+		params.EndTime = aws.Int64(until.UnixMilli())
+	}
+	attempts := 0
+
+	for {
+		output, err := cwlClient.StartQuery(context.Background(), params)
+		attempts += 1
+
+		if err != nil {
+			if errors.As(err, &rnf) && attempts <= awsservice.StandardRetries {
+				// The log group/stream hasn't been created yet, so wait and retry
+				time.Sleep(retryWaitTime)
+				continue
+			}
+
+			// if the error is not a ResourceNotFoundException, we should fail here.
+			return queryId, err
+		}
+		queryId = output.QueryId
+		return queryId, err
+	}
+}
+
+func getQueryResult(queryId *string) ([][]types.ResultField, error) {
+	attempts := 0
+	var results [][]types.ResultField
+	params := &cloudwatchlogs.GetQueryResultsInput{
+		QueryId: aws.String(*queryId),
+	}
+	for {
+		if attempts > awsservice.StandardRetries {
+			return results, errors.New("exceeded retry count")
+		}
+		result, err := cwlClient.GetQueryResults(context.Background(), params)
+		log.Printf("GetQueryResult status is: %v", result.Status)
+		attempts += 1
+		if result.Status != types.QueryStatusComplete {
+			log.Printf("GetQueryResult: sleeping for 5 seconds until status is complete")
+			time.Sleep(5 * time.Second)
+			continue
+		}
+		log.Printf("GetQueryResult: result length is %d", len(result.Results))
+		if err != nil {
+			if errors.As(err, &rnf) {
+				// The log group/stream hasn't been created yet, so wait and retry
+				time.Sleep(retryWaitTime)
+				continue
+			}
+
+			// if the error is not a ResourceNotFoundException, we should fail here.
+			return results, err
+		}
+		results = result.Results
+		return results, err
+	}
+}
+
+func getLogGroup() ([]types.LogGroup, error) {
+	attempts := 0
+	var logGroups []types.LogGroup
+	params := &cloudwatchlogs.DescribeLogGroupsInput{}
+	for {
+		output, err := cwlClient.DescribeLogGroups(context.Background(), params)
+
+		attempts += 1
+
+		if err != nil {
+			if errors.As(err, &rnf) && attempts <= awsservice.StandardRetries {
+				// The log group/stream hasn't been created yet, so wait and retry
+				time.Sleep(retryWaitTime)
+				continue
+			}
+
+			// if the error is not a ResourceNotFoundException, we should fail here.
+			return logGroups, err
+		}
+		logGroups = output.LogGroups
+		return logGroups, err
+	}
+}

test/cloudwatchlogs/resources/config_log.json

@@ -8,7 +8,7 @@
       "files": {
         "collect_list": [
           {
-            "file_path": "/tmp/cwagent_log_test.log",
+            "file_path": "/tmp/cwagent_log_test.log*",
             "log_group_name": "{instance_id}",
             "log_stream_name": "{instance_id}",
             "timezone": "UTC"