pkg/configfetcher/metadata_retriever_test.go

// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package configfetcher

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

// isValidMetadataRequest is a generic helper function to verify that the
// request matches the spec set by GCE metadata service and supports the
// behavior specified in the MetadataRetriever implementation.
func isValidMetadataRequest(r *http.Request) error {
	// Check the request is a GET request.
	if r.Method != "GET" {
		return fmt.Errorf("got a %s request, want GET request",
			r.Method)
	}

	// Check request contains the right headers.
	if _, ok := r.Header[requestHeaderKey]; !ok {
		return fmt.Errorf("missing %s header", requestHeaderKey)
	}

	return nil
}

// TestGetEntry tests that the getEntry function for the following things:
//	- makes a GET request with the required headers
//	- raises an error if the status code is not 200
//	- returns the correct response from the server
//	- returns the correct etag from the server
func TestGetEntry(t *testing.T) {
	tests := []struct {
		name           string
		entry          string
		expectResponse string
		expectErr      bool
		expectEtag     string
	}{
		{
			"InstanceCustomMetaData",
			instanceCustomDataDirectory,
			`{
				"key1": "value1",
				"key2": "value2"
			}`,
			false,
			"myEtag",
		},
		{
			"ServerError",
			instanceCustomDataDirectory,
			"",
			true,
			"",
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			// Set up a test server.
			handler := func(w http.ResponseWriter, r *http.Request) {
				if test.expectErr {
					w.WriteHeader(http.StatusInternalServerError)
					return
				}

				if err := isValidMetadataRequest(r); err != nil {
					t.Errorf("unexpected err: %v", err)
				}

				w.Header().Set("etag", test.expectEtag)

				fmt.Fprint(w, test.expectResponse)
			}
			ts := httptest.NewServer(http.HandlerFunc(handler))

			resp, etag, err := getEntry(test.entry, ts.URL)
			if err != nil {
				if !test.expectErr {
					t.Errorf("got unexpected error %v", err)
				}
			} else {
				if test.expectErr {
					t.Errorf("got %s, want error", string(resp))
				} else {
					if string(resp) != test.expectResponse {
						t.Errorf("got %s, want %s", string(resp), test.expectResponse)
					}
					if etag != test.expectEtag {
						t.Errorf("got etag %s, want %s", etag, test.expectEtag)
					}
				}
			}
		})
	}
}

// TestFetchMetadata tests that the retriever implementation correctly fetches
// the content of instance custom metadata directory.
func TestFetchMetadata(t *testing.T) {
	// Dummy key-value pairs.
	dummyMap := map[string]string{
		"key1": "val1",
		"key2": "val2",
		"key3": "\t val3 \n  ",
	}

	tests := []struct {
		name           string
		lastEtag       string
		expectResponse string
		expectEtag     string
		expectErr      bool
	}{
		{
			"InstanceCustomMetadataWait",
			"lastEtag",
			dummyMap["key1"],
			"myEtag1",
			false,
		},
		{
			"InstanceCustomMetadataNoWait",
			"",
			dummyMap["key2"],
			"myEtag2",
			false,
		},
		{
			"InstanceCustomMetadataWhitespaces",
			"",
			dummyMap["key3"],
			"myEtag3",
			false,
		},
		{
			"InstanceCustomDataServerError",
			"",
			"",
			"",
			true,
		},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			// Set up a test server.
			handler := func(w http.ResponseWriter, r *http.Request) {
				if test.expectErr {
					w.WriteHeader(http.StatusInternalServerError)
					return
				}

				if err := isValidMetadataRequest(r); err != nil {
					t.Errorf("Test %s: %v", test.name, err)
				}

				// Make sure the request is for the directory.
				if r.URL.Path != instanceCustomDataDirectory {
					t.Errorf("got URL path %s, want %s",
						r.URL.Path,
						instanceCustomDataDirectory)
				}

				expectedQuery := "recursive=true"
				if test.lastEtag != "" {
					expectedQuery += "&wait_for_change=true&last_etag=" + test.lastEtag
				}
				if r.URL.RawQuery != expectedQuery {
					t.Errorf("got HTTP request query %s, want %s",
						r.URL.RawQuery,
						expectedQuery)
				}

				w.Header().Set("etag", test.expectEtag)

				fmt.Fprint(w, test.expectResponse)
			}
			ts := httptest.NewServer(http.HandlerFunc(handler))

			resp, etag, err := FetchMetadata(test.lastEtag, ts.URL)
			if err != nil {
				if !test.expectErr {
					t.Errorf("got unexpected error %v", err)
				}
			} else {
				if test.expectErr {
					t.Errorf("got %s, want error", resp)
				} else {
					if resp != strings.TrimSpace(test.expectResponse) {
						t.Errorf("got %v, want %v", resp, test.expectResponse)
					}
					if etag != test.expectEtag {
						t.Errorf("got etag %s, want %s", etag, test.expectEtag)
					}
				}
			}

		})
	}
}