fingerprint code refactor

Some code cleanup:

* Use a field for setting EC2 metadata instead of env-vars in testing;
but keep environment variables for backward compatibility reasons

* Update tests to use testify

client/fingerprint/env_aws.go

@@ -49,31 +49,22 @@ var ec2InstanceSpeedMap = map[*regexp.Regexp]int{
 // EnvAWSFingerprint is used to fingerprint AWS metadata
 type EnvAWSFingerprint struct {
 	StaticFingerprinter
+
+	// endpoint for EC2 metadata as expected by AWS SDK
+	endpoint string
+
 	logger log.Logger
 }
 
 // NewEnvAWSFingerprint is used to create a fingerprint from AWS metadata
 func NewEnvAWSFingerprint(logger log.Logger) Fingerprint {
 	f := &EnvAWSFingerprint{
-		logger: logger.Named("env_aws"),
+		logger:   logger.Named("env_aws"),
+		endpoint: strings.TrimSuffix(os.Getenv("AWS_ENV_URL"), "/meta-data/"),
 	}
 	return f
 }
 
-func ec2MetaClient(timeout time.Duration) *ec2metadata.EC2Metadata {
-	client := &http.Client{
-		Timeout:   timeout,
-		Transport: cleanhttp.DefaultTransport(),
-	}
-
-	c := aws.NewConfig().WithHTTPClient(client)
-	if endpoint := os.Getenv("AWS_ENV_URL"); endpoint != "" {
-		endpoint = strings.TrimSuffix(endpoint, "/meta-data/")
-		c = c.WithEndpoint(endpoint)
-	}
-	return ec2metadata.New(session.New(), c)
-}
-
 func (f *EnvAWSFingerprint) Fingerprint(request *FingerprintRequest, response *FingerprintResponse) error {
 	cfg := request.Config
 
@@ -84,7 +75,7 @@ func (f *EnvAWSFingerprint) Fingerprint(request *FingerprintRequest, response *F
 		timeout = 1 * time.Millisecond
 	}
 
-	ec2meta := ec2MetaClient(timeout)
+	ec2meta := ec2MetaClient(f.endpoint, timeout)
 
 	if !ec2meta.Available() {
 		return nil
@@ -199,3 +190,16 @@ func (f *EnvAWSFingerprint) linkSpeed(ec2meta *ec2metadata.EC2Metadata) int {
 
 	return netSpeed
 }
+
+func ec2MetaClient(endpoint string, timeout time.Duration) *ec2metadata.EC2Metadata {
+	client := &http.Client{
+		Timeout:   timeout,
+		Transport: cleanhttp.DefaultTransport(),
+	}
+
+	c := aws.NewConfig().WithHTTPClient(client)
+	if endpoint != "" {
+		c = c.WithEndpoint(endpoint)
+	}
+	return ec2metadata.New(session.New(), c)
+}

client/fingerprint/env_aws_test.go

@@ -5,17 +5,18 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
-	"os"
 	"testing"
 
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/helper/testlog"
 	"github.com/hashicorp/nomad/nomad/structs"
+	"github.com/stretchr/testify/require"
 )
 
 func TestEnvAWSFingerprint_nonAws(t *testing.T) {
-	os.Setenv("AWS_ENV_URL", "http://127.0.0.1/latest/meta-data/")
 	f := NewEnvAWSFingerprint(testlog.HCLogger(t))
+	f.(*EnvAWSFingerprint).endpoint = "http://127.0.0.1/latest"
+
 	node := &structs.Node{
 		Attributes: make(map[string]string),
 	}
@@ -23,43 +24,25 @@ func TestEnvAWSFingerprint_nonAws(t *testing.T) {
 	request := &FingerprintRequest{Config: &config.Config{}, Node: node}
 	var response FingerprintResponse
 	err := f.Fingerprint(request, &response)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
-
-	if len(response.Attributes) > 0 {
-		t.Fatalf("Should not apply")
-	}
+	require.NoError(t, err)
+	require.Empty(t, response.Attributes)
 }
 
 func TestEnvAWSFingerprint_aws(t *testing.T) {
+	endpoint, cleanup := startFakeEC2Metadata(t)
+	defer cleanup()
+
 	f := NewEnvAWSFingerprint(testlog.HCLogger(t))
+	f.(*EnvAWSFingerprint).endpoint = endpoint
+
 	node := &structs.Node{
 		Attributes: make(map[string]string),
 	}
 
-	// configure mock server with fixture routes, data
-	routes := routes{}
-	if err := json.Unmarshal([]byte(aws_routes), &routes); err != nil {
-		t.Fatalf("Failed to unmarshal JSON in AWS ENV test: %s", err)
-	}
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		for _, e := range routes.Endpoints {
-			if r.RequestURI == e.Uri {
-				w.Header().Set("Content-Type", e.ContentType)
-				fmt.Fprintln(w, e.Body)
-			}
-		}
-	}))
-	defer ts.Close()
-	os.Setenv("AWS_ENV_URL", ts.URL+"/latest/meta-data/")
-
 	request := &FingerprintRequest{Config: &config.Config{}, Node: node}
 	var response FingerprintResponse
 	err := f.Fingerprint(request, &response)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
+	require.NoError(t, err)
 
 	keys := []string{
 		"platform.aws.ami-id",
@@ -78,9 +61,7 @@ func TestEnvAWSFingerprint_aws(t *testing.T) {
 		assertNodeAttributeContains(t, response.Attributes, k)
 	}
 
-	if len(response.Links) == 0 {
-		t.Fatalf("Empty links for Node in AWS Fingerprint test")
-	}
+	require.NotEmpty(t, response.Links)
 
 	// confirm we have at least instance-id and ami-id
 	for _, k := range []string{"aws.ec2"} {
@@ -88,6 +69,115 @@ func TestEnvAWSFingerprint_aws(t *testing.T) {
 	}
 }
 
+func TestNetworkFingerprint_AWS(t *testing.T) {
+	endpoint, cleanup := startFakeEC2Metadata(t)
+	defer cleanup()
+
+	f := NewEnvAWSFingerprint(testlog.HCLogger(t))
+	f.(*EnvAWSFingerprint).endpoint = endpoint
+
+	node := &structs.Node{
+		Attributes: make(map[string]string),
+	}
+
+	request := &FingerprintRequest{Config: &config.Config{}, Node: node}
+	var response FingerprintResponse
+	err := f.Fingerprint(request, &response)
+	require.NoError(t, err)
+
+	assertNodeAttributeContains(t, response.Attributes, "unique.network.ip-address")
+
+	require.NotNil(t, response.NodeResources)
+	require.Len(t, response.NodeResources.Networks, 1)
+
+	// Test at least the first Network Resource
+	net := response.NodeResources.Networks[0]
+	require.NotEmpty(t, net.IP, "Expected Network Resource to have an IP")
+	require.NotEmpty(t, net.CIDR, "Expected Network Resource to have a CIDR")
+	require.NotEmpty(t, net.Device, "Expected Network Resource to have a Device Name")
+}
+
+func TestNetworkFingerprint_AWS_network(t *testing.T) {
+	endpoint, cleanup := startFakeEC2Metadata(t)
+	defer cleanup()
+
+	f := NewEnvAWSFingerprint(testlog.HCLogger(t))
+	f.(*EnvAWSFingerprint).endpoint = endpoint
+
+	{
+		node := &structs.Node{
+			Attributes: make(map[string]string),
+		}
+
+		request := &FingerprintRequest{Config: &config.Config{}, Node: node}
+		var response FingerprintResponse
+		err := f.Fingerprint(request, &response)
+		require.NoError(t, err)
+
+		require.True(t, response.Detected, "expected response to be applicable")
+
+		assertNodeAttributeContains(t, response.Attributes, "unique.network.ip-address")
+
+		require.NotNil(t, response.NodeResources)
+		require.Len(t, response.NodeResources.Networks, 1)
+
+		// Test at least the first Network Resource
+		net := response.NodeResources.Networks[0]
+		require.NotEmpty(t, net.IP, "Expected Network Resource to have an IP")
+		require.NotEmpty(t, net.CIDR, "Expected Network Resource to have a CIDR")
+		require.NotEmpty(t, net.Device, "Expected Network Resource to have a Device Name")
+		require.Equal(t, 1000, net.MBits)
+	}
+
+	// Try again this time setting a network speed in the config
+	{
+		node := &structs.Node{
+			Attributes: make(map[string]string),
+		}
+
+		cfg := &config.Config{
+			NetworkSpeed: 10,
+		}
+
+		request := &FingerprintRequest{Config: cfg, Node: node}
+		var response FingerprintResponse
+		err := f.Fingerprint(request, &response)
+		require.NoError(t, err)
+
+		assertNodeAttributeContains(t, response.Attributes, "unique.network.ip-address")
+
+		require.NotNil(t, response.NodeResources)
+		require.Len(t, response.NodeResources.Networks, 1)
+
+		// Test at least the first Network Resource
+		net := response.NodeResources.Networks[0]
+		require.NotEmpty(t, net.IP, "Expected Network Resource to have an IP")
+		require.NotEmpty(t, net.CIDR, "Expected Network Resource to have a CIDR")
+		require.NotEmpty(t, net.Device, "Expected Network Resource to have a Device Name")
+		require.Equal(t, 10, net.MBits)
+	}
+}
+
+/// Utility functions for tests
+
+func startFakeEC2Metadata(t *testing.T) (endpoint string, cleanup func()) {
+	routes := routes{}
+	if err := json.Unmarshal([]byte(aws_routes), &routes); err != nil {
+		t.Fatalf("Failed to unmarshal JSON in AWS ENV test: %s", err)
+	}
+
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		for _, e := range routes.Endpoints {
+			if r.RequestURI == e.Uri {
+				w.Header().Set("Content-Type", e.ContentType)
+				fmt.Fprintln(w, e.Body)
+			}
+		}
+	}))
+
+	return ts.URL + "/latest", ts.Close
+}
+
 type routes struct {
 	Endpoints []*endpoint `json:"endpoints"`
 }
@@ -148,168 +238,3 @@ const aws_routes = `
   ]
 }
 `
-
-func TestNetworkFingerprint_AWS(t *testing.T) {
-	// configure mock server with fixture routes, data
-	routes := routes{}
-	if err := json.Unmarshal([]byte(aws_routes), &routes); err != nil {
-		t.Fatalf("Failed to unmarshal JSON in AWS ENV test: %s", err)
-	}
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		for _, e := range routes.Endpoints {
-			if r.RequestURI == e.Uri {
-				w.Header().Set("Content-Type", e.ContentType)
-				fmt.Fprintln(w, e.Body)
-			}
-		}
-	}))
-
-	defer ts.Close()
-	os.Setenv("AWS_ENV_URL", ts.URL+"/latest/meta-data/")
-
-	f := NewEnvAWSFingerprint(testlog.HCLogger(t))
-	node := &structs.Node{
-		Attributes: make(map[string]string),
-	}
-
-	request := &FingerprintRequest{Config: &config.Config{}, Node: node}
-	var response FingerprintResponse
-	err := f.Fingerprint(request, &response)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
-
-	assertNodeAttributeContains(t, response.Attributes, "unique.network.ip-address")
-
-	if response.NodeResources == nil || len(response.NodeResources.Networks) == 0 {
-		t.Fatal("Expected to find Network Resources")
-	}
-
-	// Test at least the first Network Resource
-	net := response.NodeResources.Networks[0]
-	if net.IP == "" {
-		t.Fatal("Expected Network Resource to have an IP")
-	}
-	if net.CIDR == "" {
-		t.Fatal("Expected Network Resource to have a CIDR")
-	}
-	if net.Device == "" {
-		t.Fatal("Expected Network Resource to have a Device Name")
-	}
-}
-
-func TestNetworkFingerprint_AWS_network(t *testing.T) {
-	// configure mock server with fixture routes, data
-	routes := routes{}
-	if err := json.Unmarshal([]byte(aws_routes), &routes); err != nil {
-		t.Fatalf("Failed to unmarshal JSON in AWS ENV test: %s", err)
-	}
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		for _, e := range routes.Endpoints {
-			if r.RequestURI == e.Uri {
-				w.Header().Set("Content-Type", e.ContentType)
-				fmt.Fprintln(w, e.Body)
-			}
-		}
-	}))
-
-	defer ts.Close()
-	os.Setenv("AWS_ENV_URL", ts.URL+"/latest/meta-data/")
-
-	f := NewEnvAWSFingerprint(testlog.HCLogger(t))
-	{
-		node := &structs.Node{
-			Attributes: make(map[string]string),
-		}
-
-		request := &FingerprintRequest{Config: &config.Config{}, Node: node}
-		var response FingerprintResponse
-		err := f.Fingerprint(request, &response)
-		if err != nil {
-			t.Fatalf("err: %v", err)
-		}
-
-		if !response.Detected {
-			t.Fatalf("expected response to be applicable")
-		}
-
-		assertNodeAttributeContains(t, response.Attributes, "unique.network.ip-address")
-
-		if response.NodeResources == nil || len(response.NodeResources.Networks) == 0 {
-			t.Fatal("Expected to find Network Resources")
-		}
-
-		// Test at least the first Network Resource
-		net := response.NodeResources.Networks[0]
-		if net.IP == "" {
-			t.Fatal("Expected Network Resource to have an IP")
-		}
-		if net.CIDR == "" {
-			t.Fatal("Expected Network Resource to have a CIDR")
-		}
-		if net.Device == "" {
-			t.Fatal("Expected Network Resource to have a Device Name")
-		}
-		if net.MBits != 1000 {
-			t.Fatalf("Expected Network Resource to have speed %d; got %d", 1000, net.MBits)
-		}
-	}
-
-	// Try again this time setting a network speed in the config
-	{
-		node := &structs.Node{
-			Attributes: make(map[string]string),
-		}
-
-		cfg := &config.Config{
-			NetworkSpeed: 10,
-		}
-
-		request := &FingerprintRequest{Config: cfg, Node: node}
-		var response FingerprintResponse
-		err := f.Fingerprint(request, &response)
-		if err != nil {
-			t.Fatalf("err: %v", err)
-		}
-
-		assertNodeAttributeContains(t, response.Attributes, "unique.network.ip-address")
-
-		if response.NodeResources == nil || len(response.NodeResources.Networks) == 0 {
-			t.Fatal("Expected to find Network Resources")
-		}
-
-		// Test at least the first Network Resource
-		net := response.NodeResources.Networks[0]
-		if net.IP == "" {
-			t.Fatal("Expected Network Resource to have an IP")
-		}
-		if net.CIDR == "" {
-			t.Fatal("Expected Network Resource to have a CIDR")
-		}
-		if net.Device == "" {
-			t.Fatal("Expected Network Resource to have a Device Name")
-		}
-		if net.MBits != 10 {
-			t.Fatalf("Expected Network Resource to have speed %d; got %d", 10, net.MBits)
-		}
-	}
-}
-
-func TestNetworkFingerprint_notAWS(t *testing.T) {
-	os.Setenv("AWS_ENV_URL", "http://127.0.0.1/latest/meta-data/")
-	f := NewEnvAWSFingerprint(testlog.HCLogger(t))
-	node := &structs.Node{
-		Attributes: make(map[string]string),
-	}
-
-	request := &FingerprintRequest{Config: &config.Config{}, Node: node}
-	var response FingerprintResponse
-	err := f.Fingerprint(request, &response)
-	if err != nil {
-		t.Fatalf("err: %v", err)
-	}
-
-	if len(response.Attributes) > 0 {
-		t.Fatalf("Should not apply")
-	}
-}

client/fingerprint/fingerprint_test.go

@@ -7,65 +7,37 @@ import (
 
 	"github.com/hashicorp/nomad/client/config"
 	"github.com/hashicorp/nomad/nomad/structs"
+	"github.com/stretchr/testify/require"
 )
 
 func assertFingerprintOK(t *testing.T, fp Fingerprint, node *structs.Node) *FingerprintResponse {
 	request := &FingerprintRequest{Config: new(config.Config), Node: node}
 	var response FingerprintResponse
 	err := fp.Fingerprint(request, &response)
-	if err != nil {
-		t.Fatalf("Failed to fingerprint: %s", err)
-	}
+	require.NoError(t, err)
 
-	if len(response.Attributes) == 0 {
-		t.Fatalf("Failed to apply node attributes")
-	}
+	require.NotEmpty(t, response.Attributes, "Failed to apply node attributes")
 
 	return &response
 }
 
 func assertNodeAttributeContains(t *testing.T, nodeAttributes map[string]string, attribute string) {
-	if nodeAttributes == nil {
-		t.Errorf("expected an initialized map for node attributes")
-		return
-	}
+	require.NotNil(t, nodeAttributes, "expected an initialized map for node attributes")
 
-	actual, found := nodeAttributes[attribute]
-	if !found {
-		t.Errorf("Expected to find Attribute `%s`\n\n[DEBUG] %#v", attribute, nodeAttributes)
-		return
-	}
-	if actual == "" {
-		t.Errorf("Expected non-empty Attribute value for `%s`\n\n[DEBUG] %#v", attribute, nodeAttributes)
-	}
+	require.Contains(t, nodeAttributes, attribute)
+	require.NotEmpty(t, nodeAttributes[attribute])
 }
 
 func assertNodeAttributeEquals(t *testing.T, nodeAttributes map[string]string, attribute string, expected string) {
-	if nodeAttributes == nil {
-		t.Errorf("expected an initialized map for node attributes")
-		return
-	}
-	actual, found := nodeAttributes[attribute]
-	if !found {
-		t.Errorf("Expected to find Attribute `%s`; unable to check value\n\n[DEBUG] %#v", attribute, nodeAttributes)
-		return
-	}
-	if expected != actual {
-		t.Errorf("Expected `%s` Attribute to be `%s`, found `%s`\n\n[DEBUG] %#v", attribute, expected, actual, nodeAttributes)
-	}
+	require.NotNil(t, nodeAttributes, "expected an initialized map for node attributes")
+
+	require.Contains(t, nodeAttributes, attribute)
+	require.Equal(t, expected, nodeAttributes[attribute])
 }
 
 func assertNodeLinksContains(t *testing.T, nodeLinks map[string]string, link string) {
-	if nodeLinks == nil {
-		t.Errorf("expected an initialized map for node links")
-		return
-	}
-	actual, found := nodeLinks[link]
-	if !found {
-		t.Errorf("Expected to find Link `%s`\n\n[DEBUG]", link)
-		return
-	}
-	if actual == "" {
-		t.Errorf("Expected non-empty Link value for `%s`\n\n[DEBUG]", link)
-	}
+	require.NotNil(t, nodeLinks, "expected an initialized map for node links")
+
+	require.Contains(t, nodeLinks, link)
+	require.NotEmpty(t, nodeLinks[link])
 }