Rewrite metadata retriever implementation to save cpu util

CPU utilization was increasing due to repeated calls to metadata
server. This implementation of metadata retriever saves the call
to metadata server and parses the metadata string in policy manager.

BUG=b/209939594
TEST=presubmit and started a VM
RELEASE_NOTE=None

Change-Id: Ibd3d4f159826f1d8548c38d9a21c55c2dca5fdb6
diff --git a/cmd/monitor/monitor.go b/cmd/monitor/monitor.go
index 0fbd25f..d925e24 100644
--- a/cmd/monitor/monitor.go
+++ b/cmd/monitor/monitor.go
@@ -53,6 +53,10 @@
 	// to do any key management on the server side.
 	privateKeyFile = "/var/lib/devicesettings/private.key"
 
+	// GCEMetadataURL is the URL of metadata server where we can query
+	// metadata tags and their values.
+	gceMetadataURL = "http://metadata.google.internal/computeMetadata/v1"
+
 	systemctlCmd = "systemctl"
 )
 
@@ -179,7 +183,7 @@
 	policyEnforcer := policyenforcer.NewPolicyEnforcer(*systemdClient)
 
 	userConfig := make(chan *protos.InstanceConfig)
-	go configfetcher.PollUserConfig(userConfig)
+	go configfetcher.PollUserConfig(userConfig, gceMetadataURL)
 
 	for {
 		instanceConfig := <-userConfig
diff --git a/pkg/configfetcher/configfetcher.go b/pkg/configfetcher/configfetcher.go
index 7a93e98..e26e77e 100644
--- a/pkg/configfetcher/configfetcher.go
+++ b/pkg/configfetcher/configfetcher.go
@@ -15,13 +15,14 @@
 package configfetcher
 
 import (
+	"encoding/json"
 	"fmt"
 	"strconv"
 	"time"
 
 	"policy-manager/protos"
 
-	"cloud.google.com/go/compute/metadata"
+	metadataServer "cloud.google.com/go/compute/metadata"
 	"github.com/golang/glog"
 	"github.com/golang/protobuf/jsonpb"
 	"github.com/golang/protobuf/proto"
@@ -52,28 +53,20 @@
 	retryInterval = 5 * time.Second
 )
 
-// FetchUserConfig returns the InstanceConfig specified by user in GCE metadata. A nil is
-// returned if there is any error.
-func FetchUserConfig() (*protos.InstanceConfig, error) {
-	if userConfig, err := getUserConfig(); err == nil {
-		return userConfig, nil
-	} else {
-		return nil, err
-	}
-}
-
 // PollUserConfig watches the update of the InstanceConfig specified by user in GCE
 // metadata. It repeatedly sends request to wait for update of userConfig. When an update
 // is returned, PollUserConfig compares it to the cached userConfig and outputs to the
 // channel if it is really updated.
-func PollUserConfig(out chan *protos.InstanceConfig) {
+func PollUserConfig(out chan *protos.InstanceConfig, gceMetadataURL string) {
 	var userConfig, lastUserConfig *protos.InstanceConfig
 	var err error
+	etag := ""
 	lastUserConfig = nil
 	for {
-		userConfig, err = getUserConfig()
+		userConfig, etag, err = getUserConfig(etag, gceMetadataURL)
 		if err != nil {
 			time.Sleep(retryInterval)
+			etag = ""
 			continue // Retry forever.
 		}
 		// Check whether userConfig is updated.
@@ -91,23 +84,40 @@
 
 // getUserConfig gets the content of the instance custom metadata keys'.
 // Then getUserConfig parses the metadata keys and returns userConfig.
-func getUserConfig() (*protos.InstanceConfig, error) {
-	var metadataKeys []string
-	var err error
-	if metadataKeys, err = metadata.InstanceAttributes(); err != nil {
-		return nil, fmt.Errorf("failed to fetch metadata from metadata server: %s", err)
-	}
-	glog.V(4).Infof("Current value of instance custom metadata keys: %v", metadataKeys)
-	userConfig, err := getUserConfigFromMetadata(metadataKeys)
+func getUserConfig(lastEtag string, gceMetadataURL string) (*protos.InstanceConfig, string, error) {
+	rawMetadata, etag, err := FetchMetadata(lastEtag, gceMetadataURL)
 	if err != nil {
+		glog.Errorf("Failed to fetch metadata from metadata server: %s", err)
+		return nil, "", err
+	}
+
+	metadata, err := parseRawMetadata(rawMetadata)
+	if err != nil {
+		glog.Errorf("Failed to parse raw metadata: %s", err)
+		return nil, "", err
+	}
+	glog.Infof("Current value of instance custom metadata keys: %v", metadata)
+	userConfig, err := getUserConfigFromMetadata(metadata)
+	if err != nil {
+		glog.Errorf("Failed to resolve userconfig: %s", err)
+		return nil, "", err
+	}
+	return userConfig, etag, nil
+}
+
+// Converts instance custom metadata from string to key:value pairs.
+func parseRawMetadata(rawMetadata string) (map[string]string, error) {
+	output := make(map[string]string)
+
+	if err := json.Unmarshal([]byte(rawMetadata), &output); err != nil {
 		return nil, err
 	}
-	return userConfig, nil
+	return output, nil
 }
 
 // GetInstanceID returns the id of the instance.
 func GetInstanceID() (uint64, error) {
-	var resp, err = metadata.InstanceID()
+	var resp, err = metadataServer.InstanceID()
 	if err != nil {
 		return 0, err
 	}
@@ -119,27 +129,17 @@
 // corresponding key with 'gci-' prefix. For example, 'cos-update-strategy' and
 // 'gci-update-strategy'. The value of 'gci-' key is returned only if the function
 // fail to get the value of 'cos-' key.
-func getCOSGCIConfigSetting(cosKey string, gciKey string, metadataKeys []string) (string, bool) {
-	var val string
-	var err error
-
-	for _, metadataKey := range metadataKeys {
-		if metadataKey == cosKey {
-			if val, err = metadata.InstanceAttributeValue(cosKey); err == nil {
-				return val, true
-			}
-		} else if metadataKey == gciKey {
-			if val, err = metadata.InstanceAttributeValue(gciKey); err == nil {
-				return val, true
-			}
-		}
+func getCOSGCIConfigSetting(cosKey string, gciKey string, metadata map[string]string) (string, bool) {
+	value, ok := metadata[cosKey]
+	if !ok {
+		value, ok = metadata[gciKey]
 	}
-	return "", false
+	return value, ok
 }
 
 // Gets user config based on the metadata keys. An empty InstanceConfig is returned
 // if there is no InstanceConfig-related metadata in the pairs.
-func getUserConfigFromMetadata(metadataKeys []string) (*protos.InstanceConfig, error) {
+func getUserConfigFromMetadata(metadata map[string]string) (*protos.InstanceConfig, error) {
 	var configStr string
 	var err error
 	var ok, boolean bool
@@ -148,7 +148,7 @@
 
 	// Get legacy instance config from metadata. This is deprecated, but if this is still
 	// being specified, use it and ignore others.
-	if configStr, ok = getCOSGCIConfigSetting("", gciLegacyConfigKey, metadataKeys); ok {
+	if configStr, ok = getCOSGCIConfigSetting("", gciLegacyConfigKey, metadata); ok {
 		if configStr != "" {
 			if err := jsonpb.UnmarshalString(configStr, userConfig); err != nil {
 				return nil, fmt.Errorf("failed to unmarshal InstanceConfig %s: %s", configStr, err)
@@ -158,7 +158,7 @@
 	}
 
 	// Get individual config settings. Keys with 'cos-' prefix have the priority.
-	if configStr, ok = getCOSGCIConfigSetting(cosKeyUpdateStrategy, gciKeyUpdateStrategy, metadataKeys); ok {
+	if configStr, ok = getCOSGCIConfigSetting(cosKeyUpdateStrategy, gciKeyUpdateStrategy, metadata); ok {
 		if configStr == "update_disabled" {
 			userConfig.UpdateStrategy = proto.String(configStr)
 		} else {
@@ -166,19 +166,19 @@
 		}
 	}
 
-	if configStr, ok = getCOSGCIConfigSetting(cosKeyMetricsEnabled, gciKeyMetricsEnabled, metadataKeys); ok {
+	if configStr, ok = getCOSGCIConfigSetting(cosKeyMetricsEnabled, gciKeyMetricsEnabled, metadata); ok {
 		if boolean, err = strconv.ParseBool(configStr); err == nil {
 			userConfig.MetricsEnabled = proto.Bool(boolean)
 		}
 	}
 
-	if configStr, ok := getCOSGCIConfigSetting(keyGoogleLoggingEnabled, "", metadataKeys); ok {
+	if configStr, ok := getCOSGCIConfigSetting(keyGoogleLoggingEnabled, "", metadata); ok {
 		if boolean, err = strconv.ParseBool(configStr); err == nil {
 			userConfig.HealthMonitorConfig.LoggingEnabled = proto.Bool(boolean)
 		}
 	}
 
-	if configStr, ok := getCOSGCIConfigSetting(keyGoogleMonitoringEnabled, "", metadataKeys); ok {
+	if configStr, ok := getCOSGCIConfigSetting(keyGoogleMonitoringEnabled, "", metadata); ok {
 		if boolean, err = strconv.ParseBool(configStr); err == nil {
 			userConfig.HealthMonitorConfig.MonitoringEnabled = proto.Bool(boolean)
 		}
diff --git a/pkg/configfetcher/configfetcher_test.go b/pkg/configfetcher/configfetcher_test.go
index dcd70e8..c8dc5cb 100644
--- a/pkg/configfetcher/configfetcher_test.go
+++ b/pkg/configfetcher/configfetcher_test.go
@@ -15,10 +15,12 @@
 package configfetcher
 
 import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
 	"testing"
 
 	"github.com/golang/protobuf/proto"
-	"policy-manager/pkg/fakes"
 	"policy-manager/protos"
 )
 
@@ -133,15 +135,7 @@
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			fakeMetadata := fakes.NewMetadataServer()
-			defer fakeMetadata.Server.Close()
-			metadataKeys := make([]string, 0, len(test.metadata))
-			for k := range test.metadata {
-				metadataKeys = append(metadataKeys, k)
-			}
-
-			fakeMetadata.Metadata = test.metadata
-			actualConfig, _ := getUserConfigFromMetadata(metadataKeys)
+			actualConfig, _ := getUserConfigFromMetadata(test.metadata)
 			if !proto.Equal(actualConfig, test.expectedConfig) {
 				t.Errorf("FAILED '%s': got %s, expect %s",
 					test.name,
@@ -159,23 +153,23 @@
 		// Name of the test case
 		name string
 		// The metadata returned by metatdata client
-		metadata map[string]string
+		metadata string
 		// The expected config settings to be returned
 		expectedConfig *protos.InstanceConfig
 	}{
 		{
 			"NoMetadata",
-			map[string]string{},
+			"{}",
 			&protos.InstanceConfig{
 				HealthMonitorConfig: &protos.HealthMonitorConfig{},
 			},
 		},
 		{
 			"NormalUpdate",
-			map[string]string{
+			`{
 				"cos-metrics-enabled": "true",
-				"cos-update-strategy": "",
-			},
+				"cos-update-strategy": ""
+			}`,
 			&protos.InstanceConfig{
 				UpdateStrategy:      proto.String(""),
 				MetricsEnabled:      proto.Bool(true),
@@ -184,11 +178,11 @@
 		},
 		{
 			"MixedUpdate",
-			map[string]string{
+			`{
 				"cos-update-strategy": "update_disabled",
 				"key1":                "value1",
-				"kdy2":                "value2",
-			},
+				"kdy2":                "value2"
+			}`,
 			&protos.InstanceConfig{
 				UpdateStrategy:      proto.String("update_disabled"),
 				HealthMonitorConfig: &protos.HealthMonitorConfig{},
@@ -196,22 +190,22 @@
 		},
 		{
 			"NoCosMetadataUpdate",
-			map[string]string{
+			`{
 				"key1": "value1",
-				"kdy2": "value2",
-			},
+				"kdy2": "value2"
+			}`,
 			&protos.InstanceConfig{
 				HealthMonitorConfig: &protos.HealthMonitorConfig{},
 			},
 		},
 		{
 			"LoggingButNoMonitoring",
-			map[string]string{
+			`{
 				"google-logging-enabled":    "true",
 				"google-monitoring-enabled": "false",
 				"key1":                      "value1",
-				"kdy2":                      "value2",
-			},
+				"kdy2":                      "value2"
+			}`,
 			&protos.InstanceConfig{
 				HealthMonitorConfig: &protos.HealthMonitorConfig{
 					Enforced:          proto.Bool(true),
@@ -222,11 +216,11 @@
 		},
 		{
 			"MonitoringButIgnoreLogging",
-			map[string]string{
+			`{
 				"google-monitoring-enabled": "true",
 				"key1":                      "value1",
-				"kdy2":                      "value2",
-			},
+				"kdy2":                      "value2"
+			}`,
 			&protos.InstanceConfig{
 				HealthMonitorConfig: &protos.HealthMonitorConfig{
 					Enforced:          proto.Bool(true),
@@ -236,12 +230,12 @@
 		},
 		{
 			"LoggingAndMonitoring",
-			map[string]string{
+			`{
 				"google-logging-enabled":    "true",
 				"google-monitoring-enabled": "true",
 				"key1":                      "value1",
-				"kdy2":                      "value2",
-			},
+				"kdy2":                      "value2"
+			}`,
 			&protos.InstanceConfig{
 				HealthMonitorConfig: &protos.HealthMonitorConfig{
 					Enforced:          proto.Bool(true),
@@ -254,10 +248,29 @@
 
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			fakeMetadata := fakes.NewMetadataServer()
-			defer fakeMetadata.Server.Close()
-			fakeMetadata.Metadata = test.metadata
-			actualConfig, _ := getUserConfig()
+			// Set up a test server.
+			handler := func(w http.ResponseWriter, r *http.Request) {
+				// 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 r.URL.RawQuery != expectedQuery {
+					t.Errorf("got HTTP request query %s, want %s",
+						r.URL.RawQuery,
+						expectedQuery)
+				}
+
+				w.Header().Set("etag", "")
+
+				fmt.Fprint(w, test.metadata)
+			}
+			ts := httptest.NewServer(http.HandlerFunc(handler))
+
+			actualConfig, _, _ := getUserConfig("", ts.URL)
 
 			if !proto.Equal(actualConfig, test.expectedConfig) {
 				t.Errorf("FAILED '%s': got %s, expect %s",
diff --git a/pkg/configfetcher/metadata_retriever.go b/pkg/configfetcher/metadata_retriever.go
new file mode 100644
index 0000000..b508afd
--- /dev/null
+++ b/pkg/configfetcher/metadata_retriever.go
@@ -0,0 +1,117 @@
+// 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
+
+// Implementation of the MetadataRetriever interface for running on an actual
+// GCE instance. This implementation allows us to reduce the number of calls
+// to metadata server, thus reducing the cpu utilization of device_policy_manager.
+// If we were using go pkg for metadata server(https://pkg.go.dev/cloud.google.com/go/compute/metadata),
+// we can use two functions: InstanceAttributes() and  InstanceAttributeValue().
+// This would have resulted in multiple calls to metadata server.
+// Current implementation allows us to call metadata server once,
+// and then parse the raw metadata on our side, thus reducing the cpu
+// utilization.
+
+// API for querying the metadata server is specified here:
+// https://cloud.google.com/compute/docs/metadata
+
+import (
+	"fmt"
+	"io/ioutil"
+	"net/http"
+	"strings"
+)
+
+// Non-exported constants for querying the metadata server.
+const (
+	// Header to include in every request to the metadata server.
+	requestHeaderKey   = "Metadata-Flavor"
+	requestHeaderValue = "Google"
+
+	// Note that there are 2 types of metadata entries: endpoint and
+	// directory. Endpoint maps to a single value, while directory contains
+	// a list of key value pairs.
+	// URL path for metadata directory. Note that the trailing slash
+	// indicates that the entry is a directory instead of an endpoint.
+	// Without it the request will just be forwared by the metadata
+	// server with a 301 status code.
+	// The leading slash is required since these entries are relative to the
+	// metadata server URL of the form http://.../...
+	instanceCustomDataDirectory = "/instance/attributes/"
+	projectCustomDataDirectory  = "/project/attributes/"
+)
+
+// getEntry performs a HTTP GET request for the metadata entry specified by
+// its url path, which is relative to the metadata server's URL.
+// It returns the body of the response if and only if the status code of the
+// response is 200.
+//
+// Args:
+//	entry: the path of the metadata entry, e.g. /instance/attributes/cos-update-strategy
+//	metadataURL: the standard url to fetch metadata from.
+// Returns:
+//	resp: the response body of the HTTP GET request.
+//	etag: "ETag" field of the response head.
+//	err: error.
+func getEntry(entry, metadataURL string) ([]byte, string, error) {
+	url := metadataURL + entry
+
+	// Construct HTTP client to make the request with the necessary headers.
+	client := &http.Client{
+		// Use default policy to stop after 10 consecutive requests.
+		CheckRedirect: nil,
+	}
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		return nil, "", err
+	}
+	req.Header.Add(requestHeaderKey, requestHeaderValue)
+
+	// Make the request and read the entire response.
+	resp, err := client.Do(req)
+	if resp != nil && resp.Body != nil {
+		defer resp.Body.Close()
+	}
+	if err != nil {
+		return nil, "", err
+	}
+	if resp.StatusCode != http.StatusOK {
+		return nil, "", fmt.Errorf("request for %s returned status %d",
+			entry, resp.StatusCode)
+	}
+
+	body, err := ioutil.ReadAll(resp.Body)
+
+	return body, resp.Header.Get("etag"), err
+}
+
+// FetchMetadata fetches the content and the ETag of the instance custom
+// metadata directory. If the parameter 'lastEtag' is not empty, FetchMetadata
+// won't return until the directory is updated (compared to 'lastEtag').
+// Otherwise, it returns immediately without waiting.
+func FetchMetadata(lastEtag string, gceMetadataURL string) (
+	string, string, error) {
+	path := instanceCustomDataDirectory + "?recursive=true"
+	if lastEtag != "" {
+		path += "&wait_for_change=true&last_etag=" + lastEtag
+	}
+
+	resp, etag, err := getEntry(path, gceMetadataURL)
+	if err != nil {
+		return "", "", err
+	}
+
+	return strings.TrimSpace(string(resp)), etag, nil
+}
diff --git a/pkg/configfetcher/metadata_retriever_test.go b/pkg/configfetcher/metadata_retriever_test.go
new file mode 100644
index 0000000..010a6c9
--- /dev/null
+++ b/pkg/configfetcher/metadata_retriever_test.go
@@ -0,0 +1,216 @@
+// 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)
+					}
+				}
+			}
+
+		})
+	}
+}
diff --git a/pkg/fakes/metadata.go b/pkg/fakes/metadata.go
deleted file mode 100644
index e1d0186..0000000
--- a/pkg/fakes/metadata.go
+++ /dev/null
@@ -1,75 +0,0 @@
-// 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 fakes contains fake implementations to be used in unit tests.
-package fakes
-
-import (
-	"log"
-	"net/http"
-	"net/http/httptest"
-	"os"
-	"strings"
-)
-
-// GCEMetadata is a fake GCE implementation. It is intended to be constructed with NewMetadataServer.
-type GCEMetadata struct {
-	// MetadataKeys represents the collection of objects that exist in the fake GCS server.
-	// Keys and values are strings of the form "/instance/attributes".
-	Metadata map[string]string
-	// Server is an HTTP server that serves fake metadata requests. Requests are served using the state stored in
-	// the other struct fields.
-	Server *httptest.Server
-}
-
-// metadataKeysHandler writes all metadata keys or values to a particular metadata key upon request
-func (mtd *GCEMetadata) metadataKeysHandler(w http.ResponseWriter, r *http.Request) {
-	metadataKey := strings.TrimPrefix(r.URL.Path, "/computeMetadata/v1/instance/attributes/")
-
-	if metadataKey == "" {
-		var res string
-		for k := range mtd.Metadata {
-			res += strings.TrimSpace(k) + "\n"
-		}
-		if _, err := w.Write([]byte(res)); err != nil {
-			log.Printf("write %q failed: %v", r.URL.Path, err)
-		}
-
-	} else {
-		res := string(mtd.Metadata[metadataKey])
-		if _, err := w.Write([]byte(res)); err != nil {
-			log.Printf("write %q failed: %v", r.URL.Path, err)
-		}
-	}
-}
-
-// NewMetadataServer constructs a fake Metadata implementation. This function also sets
-// an environment variable GCE_METADATA_HOST which is required by metadata package. The
-// use of this environment variable makes it unsafe for concurrency.
-func NewMetadataServer() *GCEMetadata {
-	mtd := &GCEMetadata{
-		Metadata: make(map[string]string),
-		Server:   nil,
-	}
-
-	mux := http.NewServeMux()
-	mux.HandleFunc("/computeMetadata/v1/instance/attributes/", mtd.metadataKeysHandler)
-	mtd.Server = httptest.NewServer(mux)
-
-	// Setting metadata host url by skipping the http:// part
-	serverURL := mtd.Server.URL[7:]
-	os.Setenv("GCE_METADATA_HOST", serverURL)
-
-	return mtd
-}
diff --git a/pkg/fakes/metadata_test.go b/pkg/fakes/metadata_test.go
deleted file mode 100644
index 06640fa..0000000
--- a/pkg/fakes/metadata_test.go
+++ /dev/null
@@ -1,70 +0,0 @@
-// 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 fakes
-
-import (
-	"cloud.google.com/go/compute/metadata"
-	"fmt"
-	"testing"
-)
-
-// Tests that all metadata keys are parsed and applied correctly.
-func TestMetadataKeysHandler(t *testing.T) {
-	tests := []struct {
-		// Name of the test case
-		name string
-		// Key:Value pairs present in metadata.
-		metadata map[string]string
-		// Expected metadataKeys to be returned
-		expectedMetadataKeys []string
-	}{
-		{
-			"NoMetadataKeys",
-			map[string]string{},
-			[]string{""},
-		},
-		{
-			"UpdateStrategyPresent",
-			map[string]string{"update-strategy": "update-strategy-1"},
-			[]string{"update-strategy"},
-		},
-		{
-			"TwoKeysPresent",
-			map[string]string{"cos-update-strategy": "update-strategy-1",
-				"google-logging-enabled": "true"},
-			[]string{"cos-update-strategy", "google-logging-enabled"},
-		},
-	}
-
-	for _, test := range tests {
-		t.Run(test.name, func(t *testing.T) {
-			fakeMetadata := NewMetadataServer()
-			fakeMetadata.Metadata = test.metadata
-			defer fakeMetadata.Server.Close()
-			if res, _ := metadata.InstanceAttributes(); len(res) != len(test.expectedMetadataKeys) {
-				fmt.Printf("%v", test.expectedMetadataKeys)
-				t.Errorf("FAILED '%s': got %v, expect %v",
-					test.name, res, test.expectedMetadataKeys)
-			} else {
-				for _, val := range res {
-					if metadataVal, _ := metadata.InstanceAttributeValue(val); metadataVal != test.metadata[val] {
-						t.Errorf("FAILED '%s': got %v, expect %v",
-							test.name, metadataVal, test.metadata[val])
-					}
-				}
-			}
-		})
-	}
-}