Shutdown attribute (#855)

* guest attribute implementation for shutdown script

* print statement

* more printsA

* Fix boolean check

* Check difF

* Fix typo

imagetest/cmd/wrapper/main.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net/http"
 	"net/url"
 	"os/exec"
 	"strings"
@@ -16,19 +17,44 @@ import (
 	"github.com/GoogleCloudPlatform/guest-test-infra/imagetest/utils"
 )
 
+// In special cases such as the shutdown script, the guest attribute match
+// on the first boot must have a different name than the usual guest attribute.
+func checkFirstBootSpecialGA() bool {
+	if _, err := utils.GetMetadataAttribute("shouldRebootDuringTest"); err == nil {
+		_, foundFirstBootGA := utils.GetMetadataGuestAttribute(utils.GuestAttributeTestNamespace + "/" + utils.FirstBootGAKey)
+		// if the special attribute to match the first boot of the shutdown script test is already set, foundFirstBootGA will be nil and we should use the regular guest attribute.
+		if foundFirstBootGA != nil {
+			return true
+		}
+	}
+	return false
+}
+
 func main() {
 	ctx := context.Background()
 	client, err := storage.NewClient(ctx)
 	if err != nil {
 		log.Fatalf("failed to create cloud storage client: %v", err)
 	}
 	log.Printf("FINISHED-BOOTING")
-	defer func() {
+	firstBootSpecialAttribute := checkFirstBootSpecialGA()
+	// firstBootSpecialGA should be true if we need to match a different guest attribute than the usual guest attribute
+	defer func(ctx context.Context, firstBootSpecialGA bool) {
+		var err error
+		if firstBootSpecialGA {
+			err = utils.QueryMetadataGuestAttribute(ctx, utils.GuestAttributeTestNamespace, utils.FirstBootGAKey, http.MethodPut)
+		} else {
+			err = utils.QueryMetadataGuestAttribute(ctx, utils.GuestAttributeTestNamespace, utils.GuestAttributeTestKey, http.MethodPut)
+		}
+
+		if err != nil {
+			log.Printf("could not place guest attribute key to end test")
+		}
 		for f := 0; f < 5; f++ {
 			log.Printf("FINISHED-TEST")
 			time.Sleep(1 * time.Second)
 		}
-	}()
+	}(ctx, firstBootSpecialAttribute)
 
 	daisyOutsPath, err := utils.GetMetadataAttribute("daisy-outs-path")
 	if err != nil {

imagetest/testworkflow.go

@@ -235,7 +235,13 @@ func (t *TestWorkflow) addWaitStep(stepname, vmname string) (*daisy.Step, error)
 	instanceSignal.Name = vmname
 	instanceSignal.Stopped = false
 
+	guestAttribute := &daisy.GuestAttribute{}
+	guestAttribute.Namespace = utils.GuestAttributeTestNamespace
+	guestAttribute.KeyName = utils.GuestAttributeTestKey
+
 	instanceSignal.SerialOutput = serialOutput
+	instanceSignal.GuestAttribute = guestAttribute
+	instanceSignal.Interval = "8s"
 
 	waitForInstances := &daisy.WaitForInstancesSignal{instanceSignal}
 
@@ -258,7 +264,15 @@ func (t *TestWorkflow) addWaitRebootGAStep(stepname, vmname string) (*daisy.Step
 	instanceSignal.Name = vmname
 	instanceSignal.Stopped = false
 
+	guestAttribute := &daisy.GuestAttribute{}
+	guestAttribute.Namespace = utils.GuestAttributeTestNamespace
+	// specifically wait for a different guest attribute if this is the
+	// first boot before a reboot, and we want test results from a reboot.
+	guestAttribute.KeyName = utils.FirstBootGAKey
+
 	instanceSignal.SerialOutput = serialOutput
+	instanceSignal.GuestAttribute = guestAttribute
+	instanceSignal.Interval = "8s"
 
 	waitForInstances := &daisy.WaitForInstancesSignal{instanceSignal}
 

imagetest/utils/test_utils.go

@@ -28,8 +28,10 @@ const (
 	bytesInGB         = 1073741824
 	// GuestAttributeTestNamespace is the namespace for the guest attribute in the daisy "wait for instance" step for CIT.
 	GuestAttributeTestNamespace = "citTest"
-	// GuestAttributeTestKey is the key for the guest attribute in the edaisy "wait for instance" step for CIT.
+	// GuestAttributeTestKey is the key for the guest attribute in the daisy "wait for instance" step for CIT in the common case.
 	GuestAttributeTestKey = "test-complete"
+	// FirstBootGAKey is the key for guest attribute in the daisy "wait for instance" step in the case where it is the first boot, and we still want to wait for results from a subsequent reboot.
+	FirstBootGAKey = "first-boot-key"
 )
 
 var windowsClientImagePatterns = []string{
@@ -109,22 +111,24 @@ func GetMetadataHTTPResponse(path string) (*http.Response, error) {
 	return resp, nil
 }
 
-// PutMetadataGuestAttribute sets the guest attribute in the namespace, and returns an error if this operation fails.
-func PutMetadataGuestAttribute(ctx context.Context, namespace, attribute string) error {
+// QueryMetadataGuestAttribute sets the guest attribute in the namespace
+// using the given method, and returns an error if this operation fails.
+func QueryMetadataGuestAttribute(ctx context.Context, namespace, attribute, httpMethod string) error {
 	path, err := url.JoinPath(metadataURLPrefix, "guest-attributes", namespace, attribute)
 	if err != nil {
 		return err
 	}
-	err = PutMetadataHTTPResponse(ctx, path)
+	err = QueryMetadataHTTPResponse(ctx, path, httpMethod)
 	if err != nil {
 		return err
 	}
 	return nil
 }
 
-// PutMetadataHTTPResponse returns http response for the specified key without checking status code.
-func PutMetadataHTTPResponse(ctx context.Context, path string) error {
-	req, err := http.NewRequestWithContext(ctx, http.MethodPut, path, nil)
+// QueryMetadataHTTPResponse returns http response for the specified key
+// using a http request with the given method without checking status code.
+func QueryMetadataHTTPResponse(ctx context.Context, path, httpMethod string) error {
+	req, err := http.NewRequestWithContext(ctx, httpMethod, path, nil)
 	if err != nil {
 		return err
 	}