Check the usage_data in the config. (#938)

aws · Oct 30, 2023 · e1bdbd3 · e1bdbd3
1 parent dcbdde0
commit e1bdbd3
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 11 deletions.
diff --git a/translator/translate/otel/extension/agenthealth/translator.go b/translator/translate/otel/extension/agenthealth/translator.go
@@ -18,6 +18,8 @@ const (
 	OperationPutMetricData    = "PutMetricData"
 	OperationPutLogEvents     = "PutLogEvents"
 	OperationPutTraceSegments = "PutTraceSegments"
+
+	usageDataKey = "usage_data"
 )
 
 var (
@@ -49,9 +51,12 @@ func (t *translator) ID() component.ID {
 }
 
 // Translate creates an extension configuration.
-func (t *translator) Translate(*confmap.Conf) (component.Config, error) {
+func (t *translator) Translate(conf *confmap.Conf) (component.Config, error) {
 	cfg := t.factory.CreateDefaultConfig().(*agenthealth.Config)
 	cfg.IsUsageDataEnabled = t.isUsageDataEnabled
+	if usageData, ok := common.GetBool(conf, common.ConfigKey(common.AgentKey, usageDataKey)); ok {
+		cfg.IsUsageDataEnabled = cfg.IsUsageDataEnabled && usageData
+	}
 	cfg.Stats = agent.StatsConfig{Operations: t.operations}
 	return cfg, nil
 }
diff --git a/translator/translate/otel/extension/agenthealth/translator_test.go b/translator/translate/otel/extension/agenthealth/translator_test.go
@@ -7,21 +7,69 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"go.opentelemetry.io/collector/confmap"
 
 	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth"
 	"github.com/aws/amazon-cloudwatch-agent/extension/agenthealth/handler/stats/agent"
 )
 
 func TestTranslate(t *testing.T) {
 	operations := []string{OperationPutLogEvents}
-	tt := NewTranslator("test", operations).(*translator)
-	assert.Equal(t, "agenthealth/test", tt.ID().String())
-	tt.isUsageDataEnabled = true
-	got, err := tt.Translate(nil)
-	assert.NoError(t, err)
-	assert.Equal(t, &agenthealth.Config{IsUsageDataEnabled: true, Stats: agent.StatsConfig{Operations: operations}}, got)
-	tt.isUsageDataEnabled = false
-	got, err = tt.Translate(nil)
-	assert.NoError(t, err)
-	assert.Equal(t, &agenthealth.Config{IsUsageDataEnabled: false, Stats: agent.StatsConfig{Operations: operations}}, got)
+	testCases := map[string]struct {
+		input          map[string]interface{}
+		isEnvUsageData bool
+		want           *agenthealth.Config
+	}{
+		"WithUsageData/NotInConfig": {
+			input:          map[string]interface{}{"agent": map[string]interface{}{}},
+			isEnvUsageData: true,
+			want: &agenthealth.Config{
+				IsUsageDataEnabled: true,
+				Stats: agent.StatsConfig{
+					Operations: operations,
+				},
+			},
+		},
+		"WithUsageData/FalseInConfig": {
+			input:          map[string]interface{}{"agent": map[string]interface{}{"usage_data": false}},
+			isEnvUsageData: true,
+			want: &agenthealth.Config{
+				IsUsageDataEnabled: false,
+				Stats: agent.StatsConfig{
+					Operations: operations,
+				},
+			},
+		},
+		"WithUsageData/FalseInEnv": {
+			input:          map[string]interface{}{"agent": map[string]interface{}{"usage_data": true}},
+			isEnvUsageData: false,
+			want: &agenthealth.Config{
+				IsUsageDataEnabled: false,
+				Stats: agent.StatsConfig{
+					Operations: operations,
+				},
+			},
+		},
+		"WithUsageData/BothTrue": {
+			input:          map[string]interface{}{"agent": map[string]interface{}{"usage_data": true}},
+			isEnvUsageData: true,
+			want: &agenthealth.Config{
+				IsUsageDataEnabled: true,
+				Stats: agent.StatsConfig{
+					Operations: operations,
+				},
+			},
+		},
+	}
+	for name, testCase := range testCases {
+		t.Run(name, func(t *testing.T) {
+			tt := NewTranslator("test", operations).(*translator)
+			assert.Equal(t, "agenthealth/test", tt.ID().String())
+			tt.isUsageDataEnabled = testCase.isEnvUsageData
+			conf := confmap.NewFromStringMap(testCase.input)
+			got, err := tt.Translate(conf)
+			assert.NoError(t, err)
+			assert.Equal(t, testCase.want, got)
+		})
+	}
 }