Rewrite tests for env tags with new test helper

DataDog · Jul 30, 2021 · f7723fd · f7723fd
1 parent 39b5f93
commit f7723fd
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 152 deletions.
diff --git a/statsd/client_test.go b/statsd/client_test.go
@@ -7,7 +7,6 @@ import (
 	"net"
 	"os"
 	"reflect"
-	"sort"
 	"strings"
 	"sync"
 	"testing"
@@ -256,144 +255,6 @@ func TestNilError(t *testing.T) {
 	}
 }
 
-func TestEntityID(t *testing.T) {
-	entityIDEnvName := "DD_ENTITY_ID"
-	initialValue, initiallySet := os.LookupEnv(entityIDEnvName)
-	if initiallySet {
-		defer os.Setenv(entityIDEnvName, initialValue)
-	} else {
-		defer os.Unsetenv(entityIDEnvName)
-	}
-
-	// Set to a valid value
-	os.Setenv(entityIDEnvName, "testing")
-	client, err := statsd.New("localhost:8125")
-	if err != nil {
-		t.Fatal(err)
-	}
-	if len(client.Tags) != 1 {
-		t.Errorf("Expecting one tag, got %d", len(client.Tags))
-	}
-	if client.Tags[0] != "dd.internal.entity_id:testing" {
-		t.Errorf("Bad tag value, got %s", client.Tags[0])
-	}
-
-	// Set to empty string
-	os.Setenv(entityIDEnvName, "")
-	client, err = statsd.New("localhost:8125")
-	if err != nil {
-		t.Fatal(err)
-	}
-	if len(client.Tags) != 0 {
-		t.Errorf("Expecting empty default tags, got %v", client.Tags)
-	}
-
-	// Unset
-	os.Unsetenv(entityIDEnvName)
-	client, err = statsd.New("localhost:8125")
-	if err != nil {
-		t.Fatal(err)
-	}
-	if len(client.Tags) != 0 {
-		t.Errorf("Expecting empty default tags, got %v", client.Tags)
-	}
-}
-
-var (
-	ddEnvName     = "DD_ENV"
-	ddServiceName = "DD_SERVICE"
-	ddVersionName = "DD_VERSION"
-)
-
-func TestDDEnvServiceVersionSet(t *testing.T) {
-	for _, tt := range []struct {
-		DDEnv     string
-		DDService string
-		DDVersion string
-		Expected  []string
-	}{
-		{"", "", "", []string{}},
-		{"prod", "", "", []string{"env:prod"}},
-		{"prod", "dog", "", []string{"env:prod", "service:dog"}},
-		{"prod", "dog", "abc123", []string{"env:prod", "service:dog", "version:abc123"}},
-	} {
-		for _, t := range []string{ddEnvName, ddServiceName, ddVersionName} {
-			initialValue, initiallySet := os.LookupEnv(t)
-			if initiallySet {
-				defer os.Setenv(t, initialValue)
-			} else {
-				defer os.Unsetenv(t)
-			}
-		}
-		os.Setenv(ddEnvName, tt.DDEnv)
-		os.Setenv(ddServiceName, tt.DDService)
-		os.Setenv(ddVersionName, tt.DDVersion)
-		client, err := statsd.New("localhost:8125")
-		if err != nil {
-			t.Fatal(err)
-		}
-		// Keep the ordering of global tags consistent.
-		sort.Strings(client.Tags)
-		assert.Equal(t, tt.Expected, client.Tags)
-	}
-}
-
-func TestDDEnvServiceVersionTagsEmitted(t *testing.T) {
-	for _, t := range []string{ddEnvName, ddServiceName, ddVersionName} {
-		initialValue, initiallySet := os.LookupEnv(t)
-		if initiallySet {
-			defer os.Setenv(t, initialValue)
-		} else {
-			defer os.Unsetenv(t)
-		}
-	}
-	os.Setenv(ddEnvName, "prod")
-	os.Setenv(ddServiceName, "dog")
-	os.Setenv(ddVersionName, "abc123")
-	addr := "localhost:1201"
-	udpAddr, err := net.ResolveUDPAddr("udp", addr)
-	if err != nil {
-		t.Fatal(err)
-	}
-	server, err := net.ListenUDP("udp", udpAddr)
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer server.Close()
-
-	for _, tt := range []struct {
-		Tags       []string
-		GlobalTags []string
-		Expected   string
-	}{
-		{nil, nil, "test.count:100|c|#env:prod,service:dog,version:abc123\n"},
-		{[]string{"env:staging", "service:cat", "custom_tag"}, nil, "test.count:100|c|#env:prod,service:dog,version:abc123,env:staging,service:cat,custom_tag\n"},
-		{nil, []string{"version:def456", "custom_tag_two"}, "test.count:100|c|#custom_tag_two,env:prod,service:dog,version:abc123,version:def456\n"},
-		{[]string{"env:staging", "service:cat", "custom_tag"}, []string{"version:def456", "custom_tag_two"}, "test.count:100|c|#custom_tag_two,env:prod,service:dog,version:abc123,version:def456,env:staging,service:cat,custom_tag\n"},
-	} {
-		client, err := statsd.New(addr, statsd.WithTags(tt.GlobalTags))
-		if err != nil {
-			t.Fatal(err)
-		}
-		// Keep the ordering of global tags consistent.
-		sort.Strings(client.Tags)
-		client.Count("test.count", 100, tt.Tags, 1.0)
-		err = client.Flush()
-		if err != nil {
-			t.Errorf("Error sending: %s", err)
-		}
-		buffer := make([]byte, 1024)
-		n, err := io.ReadAtLeast(server, buffer, 1)
-		if err != nil {
-			t.Errorf("ReadAtLeast: %s", err)
-		}
-		result := string(buffer[:n])
-		if result != tt.Expected {
-			t.Errorf("Flushed metric incorrect; expected %s but got %s", tt.Expected, result)
-		}
-	}
-}
-
 func TestClosePanic(t *testing.T) {
 	c, err := statsd.New("localhost:8125")
 	assert.NoError(t, err)

diff --git a/statsd/statsd.go b/statsd/statsd.go
@@ -70,17 +70,13 @@ const (
 
 /*
 ddEnvTagsMapping is a mapping of each "DD_" prefixed environment variable
-to a specific tag name.
+to a specific tag name. We use a slice to keep the order and simplify tests.
 */
-var ddEnvTagsMapping = map[string]string{
-	// Client-side entity ID injection for container tagging.
-	"DD_ENTITY_ID": "dd.internal.entity_id",
-	// The name of the env in which the service runs.
-	"DD_ENV": "env",
-	// The name of the running service.
-	"DD_SERVICE": "service",
-	// The current version of the running service.
-	"DD_VERSION": "version",
+var ddEnvTagsMapping = []struct{ envName, tagName string }{
+	{"DD_ENTITY_ID", "dd.internal.entity_id"}, // Client-side entity ID injection for container tagging.
+	{"DD_ENV", "env"},                         // The name of the env in which the service runs.
+	{"DD_SERVICE", "service"},                 // The name of the running service.
+	{"DD_VERSION", "version"},                 // The current version of the running service.
 }
 
 type metricType int
@@ -332,9 +328,9 @@ func newWithWriter(w statsdWriter, o *Options, writerName string) (*Client, erro
 		metrics:   &ClientMetrics{},
 	}
 	// Inject values of DD_* environment variables as global tags.
-	for envName, tagName := range ddEnvTagsMapping {
-		if value := os.Getenv(envName); value != "" {
-			c.Tags = append(c.Tags, fmt.Sprintf("%s:%s", tagName, value))
+	for _, mapping := range ddEnvTagsMapping {
+		if value := os.Getenv(mapping.envName); value != "" {
+			c.Tags = append(c.Tags, fmt.Sprintf("%s:%s", mapping.tagName, value))
 		}
 	}
 

diff --git a/statsd/statsd_test.go b/statsd/statsd_test.go
@@ -5,6 +5,7 @@ import (
 	"io"
 	"net"
 	"os"
+	"sort"
 	"sync"
 	"testing"
 	"time"
@@ -333,3 +334,88 @@ func TestResolveAddressFromEnvironment(t *testing.T) {
 		})
 	}
 }
+
+func TestEnvTags(t *testing.T) {
+	entityIDEnvName := "DD_ENTITY_ID"
+	ddEnvName := "DD_ENV"
+	ddServiceName := "DD_SERVICE"
+	ddVersionName := "DD_VERSION"
+
+	defer func() { os.Unsetenv(entityIDEnvName) }()
+	defer func() { os.Unsetenv(ddEnvName) }()
+	defer func() { os.Unsetenv(ddServiceName) }()
+	defer func() { os.Unsetenv(ddVersionName) }()
+
+	os.Setenv(entityIDEnvName, "test_id")
+	os.Setenv(ddEnvName, "test_env")
+	os.Setenv(ddServiceName, "test_service")
+	os.Setenv(ddVersionName, "test_version")
+
+	expectedTags := []string{"dd.internal.entity_id:test_id", "env:test_env", "service:test_service", "version:test_version"}
+	ts, client := newClientAndTestServer(t,
+		"udp",
+		"localhost:8765",
+		expectedTags,
+	)
+
+	sort.Strings(client.Tags)
+	assert.Equal(t, client.Tags, expectedTags)
+	ts.sendAllAndAssert(t, client)
+}
+
+func TestEnvTagsWithCustomTags(t *testing.T) {
+	entityIDEnvName := "DD_ENTITY_ID"
+	ddEnvName := "DD_ENV"
+	ddServiceName := "DD_SERVICE"
+	ddVersionName := "DD_VERSION"
+
+	defer func() { os.Unsetenv(entityIDEnvName) }()
+	defer func() { os.Unsetenv(ddEnvName) }()
+	defer func() { os.Unsetenv(ddServiceName) }()
+	defer func() { os.Unsetenv(ddVersionName) }()
+
+	os.Setenv(entityIDEnvName, "test_id")
+	os.Setenv(ddEnvName, "test_env")
+	os.Setenv(ddServiceName, "test_service")
+	os.Setenv(ddVersionName, "test_version")
+
+	expectedTags := []string{"tag1", "tag2", "dd.internal.entity_id:test_id", "env:test_env", "service:test_service", "version:test_version"}
+	ts, client := newClientAndTestServer(t,
+		"udp",
+		"localhost:8765",
+		expectedTags,
+		WithTags([]string{"tag1", "tag2"}),
+	)
+
+	ts.sendAllAndAssert(t, client)
+
+	sort.Strings(expectedTags)
+	sort.Strings(client.Tags)
+	assert.Equal(t, client.Tags, expectedTags)
+}
+
+func TestEnvTagsEmptyString(t *testing.T) {
+	entityIDEnvName := "DD_ENTITY_ID"
+	ddEnvName := "DD_ENV"
+	ddServiceName := "DD_SERVICE"
+	ddVersionName := "DD_VERSION"
+
+	defer func() { os.Unsetenv(entityIDEnvName) }()
+	defer func() { os.Unsetenv(ddEnvName) }()
+	defer func() { os.Unsetenv(ddServiceName) }()
+	defer func() { os.Unsetenv(ddVersionName) }()
+
+	os.Setenv(entityIDEnvName, "")
+	os.Setenv(ddEnvName, "")
+	os.Setenv(ddServiceName, "")
+	os.Setenv(ddVersionName, "")
+
+	ts, client := newClientAndTestServer(t,
+		"udp",
+		"localhost:8765",
+		nil,
+	)
+
+	assert.Len(t, client.Tags, 0)
+	ts.sendAllAndAssert(t, client)
+}