diff --git a/imagetest/cmd/wrapper/main.go b/imagetest/cmd/wrapper/main.go index f2349b141..e9b81eeb2 100644 --- a/imagetest/cmd/wrapper/main.go +++ b/imagetest/cmd/wrapper/main.go @@ -7,6 +7,7 @@ import ( "io" "io/ioutil" "log" + "net/http" "net/url" "os/exec" "strings" @@ -16,6 +17,19 @@ 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) @@ -23,12 +37,24 @@ func main() { 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 { diff --git a/imagetest/testworkflow.go b/imagetest/testworkflow.go index ef96bc6c2..322dd5356 100644 --- a/imagetest/testworkflow.go +++ b/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} diff --git a/imagetest/utils/test_utils.go b/imagetest/utils/test_utils.go index cbbfd26f4..ad9296736 100644 --- a/imagetest/utils/test_utils.go +++ b/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 }