From a54c0de5cd7290b97e62a91ec91b443cfe25d103 Mon Sep 17 00:00:00 2001
From: Anthony Lu <67125539+koln67@users.noreply.github.com>
Date: Tue, 10 Oct 2023 13:51:21 -0700
Subject: [PATCH] Shutdown Script Minor Change (#848)

* Shutdown reboot special step

* gofmt

* Remove unused print

* Reuse current method name for testvm
---
 imagetest/fixtures.go                   | 13 +++++++++++-
 imagetest/test_suites/metadata/setup.go | 10 +++++----
 imagetest/testworkflow.go               | 27 +++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/imagetest/fixtures.go b/imagetest/fixtures.go
index c2388bb92..b68052a17 100644
--- a/imagetest/fixtures.go
+++ b/imagetest/fixtures.go
@@ -34,6 +34,9 @@ const (
 	createFirewallStepName   = "create-firewalls"
 	createSubnetworkStepName = "create-sub-networks"
 	successMatch             = "FINISHED-TEST"
+	// ShouldRebootDuringTest is a local map key to indicate that the
+	// test will reboot and relies on results from the second boot.
+	ShouldRebootDuringTest = "shouldRebootDuringTest"
 	// DefaultSourceRange is the RFC-1918 range used in default rules.
 	DefaultSourceRange = "10.128.0.0/9"
 
@@ -181,7 +184,15 @@ func (t *TestWorkflow) CreateTestVMMultipleDisks(disks []*compute.Disk, instance
 		}
 	}
 
-	waitStep, err := t.addWaitStep(vmname, vmname)
+	// In a follow-up, guest attribute support will be added.
+	// If this is the first boot before a reboot, this should use a
+	// different guest attribute when waiting for the instance signal.
+	var waitStep *daisy.Step
+	if _, foundKey := instanceParams[ShouldRebootDuringTest]; foundKey {
+		waitStep, err = t.addWaitRebootGAStep(vmname, vmname)
+	} else {
+		waitStep, err = t.addWaitStep(vmname, vmname)
+	}
 	if err != nil {
 		return nil, err
 	}
diff --git a/imagetest/test_suites/metadata/setup.go b/imagetest/test_suites/metadata/setup.go
index db8cf77c6..9fbce287a 100644
--- a/imagetest/test_suites/metadata/setup.go
+++ b/imagetest/test_suites/metadata/setup.go
@@ -5,6 +5,7 @@ import (
 	"strings"
 
 	"github.com/GoogleCloudPlatform/guest-test-infra/imagetest"
+	"google.golang.org/api/compute/v1"
 )
 
 // Name is the name of the test package. It must match the directory name.
@@ -37,7 +38,8 @@ func TestSetup(t *imagetest.TestWorkflow) error {
 		return err
 	}
 
-	vm2, err := t.CreateTestVM("vm2")
+	shutdownScriptParams := map[string]string{imagetest.ShouldRebootDuringTest: "true"}
+	vm2, err := t.CreateTestVMMultipleDisks([]*compute.Disk{{Name: "vm2"}}, shutdownScriptParams)
 	if err != nil {
 		return err
 	}
@@ -46,7 +48,7 @@ func TestSetup(t *imagetest.TestWorkflow) error {
 		return err
 	}
 
-	vm3, err := t.CreateTestVM("vm3")
+	vm3, err := t.CreateTestVMMultipleDisks([]*compute.Disk{{Name: "vm3"}}, shutdownScriptParams)
 	if err != nil {
 		return err
 	}
@@ -55,7 +57,7 @@ func TestSetup(t *imagetest.TestWorkflow) error {
 		return err
 	}
 
-	vm4, err := t.CreateTestVM("vm4")
+	vm4, err := t.CreateTestVMMultipleDisks([]*compute.Disk{{Name: "vm4"}}, shutdownScriptParams)
 	if err != nil {
 		return err
 	}
@@ -64,7 +66,7 @@ func TestSetup(t *imagetest.TestWorkflow) error {
 		return err
 	}
 
-	vm5, err := t.CreateTestVM("vm5")
+	vm5, err := t.CreateTestVMMultipleDisks([]*compute.Disk{{Name: "vm5"}}, shutdownScriptParams)
 	if err != nil {
 		return err
 	}
diff --git a/imagetest/testworkflow.go b/imagetest/testworkflow.go
index d5ac21e7f..ef96bc6c2 100644
--- a/imagetest/testworkflow.go
+++ b/imagetest/testworkflow.go
@@ -122,6 +122,10 @@ func (t *TestWorkflow) appendCreateVMStep(disks []*compute.Disk, instanceParams
 	}
 
 	instance.Metadata = make(map[string]string)
+	// set this metadata key to indicate to the wrapper that the we are running a test with both a boot and a reboot.
+	if shouldRebootDuringTest, foundKey := instanceParams[ShouldRebootDuringTest]; foundKey {
+		instance.Metadata[ShouldRebootDuringTest] = shouldRebootDuringTest
+	}
 	instance.Metadata["_test_vmname"] = name
 	instance.Metadata["_test_package_url"] = "${SOURCESPATH}/testpackage"
 	instance.Metadata["_test_results_url"] = fmt.Sprintf("${OUTSPATH}/%s.txt", name)
@@ -244,6 +248,29 @@ func (t *TestWorkflow) addWaitStep(stepname, vmname string) (*daisy.Step, error)
 	return waitStep, nil
 }
 
+// after guest attributes for instance wait step matching are implemented, this step will wait for a different guest attribute key than addWaitStep
+func (t *TestWorkflow) addWaitRebootGAStep(stepname, vmname string) (*daisy.Step, error) {
+	serialOutput := &daisy.SerialOutput{}
+	serialOutput.Port = 1
+	serialOutput.SuccessMatch = successMatch
+
+	instanceSignal := &daisy.InstanceSignal{}
+	instanceSignal.Name = vmname
+	instanceSignal.Stopped = false
+
+	instanceSignal.SerialOutput = serialOutput
+
+	waitForInstances := &daisy.WaitForInstancesSignal{instanceSignal}
+
+	waitStep, err := t.wf.NewStep("wait-" + stepname)
+	if err != nil {
+		return nil, err
+	}
+	waitStep.WaitForInstancesSignal = waitForInstances
+
+	return waitStep, nil
+}
+
 func (t *TestWorkflow) addStopStep(stepname, vmname string) (*daisy.Step, error) {
 	stopInstances := &daisy.StopInstances{}
 	stopInstances.Instances = append(stopInstances.Instances, vmname)