github.com/GoogleCloudPlatform/compute-image-tools/cli_tools@v0.0.0-20240516224744-de2dabc4ed1b/gce_windows_upgrade/upgrader/workflows_test.go

//  Copyright 2020 Google Inc. All Rights Reserved.
//
//  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 upgrader

import (
	"testing"

	"github.com/stretchr/testify/assert"

	daisy "github.com/GoogleCloudPlatform/compute-daisy"

	"github.com/GoogleCloudPlatform/compute-image-tools/cli_tools/common/utils/daisyutils"
)

func init() {
	initTest()
}

func TestGeneratePrepareWorkflow(t *testing.T) {
	type testCase struct {
		testName            string
		populateFunc        func(*upgrader, *daisy.Workflow) error
		instanceName        string
		createMachineBackup bool
	}

	tcs := []testCase{
		{"prepare", populatePrepareSteps, testInstance, false},
		{"prepare with original startup script", populatePrepareSteps, testInstanceWithStartupScript, true},
	}

	for _, tc := range tcs {
		u := newTestUpgrader()
		u.Instance = daisyutils.GetInstanceURI(testProject, testZone, tc.instanceName)

		err := u.validateAndDeriveParams()
		if err != nil {
			t.Errorf("[%v]: validateAndDeriveParams failed: %v", tc.testName, err)
			continue
		}

		w, err := u.generateWorkflowWithSteps("test", DefaultTimeout, tc.populateFunc)
		assert.NoError(t, err, "[test name: %v] Unexpected error.", tc.testName)

		_, hasBackupMachineImageStep := w.Steps["backup-machine-image"]
		if !u.CreateMachineBackup && hasBackupMachineImageStep {
			t.Errorf("[%v]: Skiped machine image backup but still see this step in workflow.", tc.testName)
		} else if u.CreateMachineBackup && !hasBackupMachineImageStep {
			t.Errorf("[%v]: Didn't skip machine image backup but can't see this step in workflow.", tc.testName)
		}

		_, hasBackupStartupScriptStep := w.Steps["backup-script"]
		if tc.instanceName != testInstanceWithStartupScript && hasBackupStartupScriptStep {
			t.Errorf("[%v]: Original startup script doesn't exist but still see this step in workflow.", tc.testName)
		} else if tc.instanceName == testInstanceWithStartupScript && !hasBackupStartupScriptStep {
			t.Errorf("[%v]: Original startup script exists but can't see this step in workflow.", tc.testName)
		}
	}
}

func TestGenerateStaticWorkflow(t *testing.T) {
	type testCase struct {
		testName     string
		populateFunc func(*upgrader, *daisy.Workflow) error
		instanceName string
	}

	tcs := []testCase{
		{"reboot", populateRebootSteps, testInstance},
		{"cleanup", populateCleanupSteps, testInstance},
		{"rollback", populateRollbackSteps, testInstanceWithStartupScript},
	}
	for _, targetOS := range SupportedVersions {
		// 2008r2 can't be target OS, skip it.
		if targetOS != versionWindows2008r2 {
			tcs = append(tcs, testCase{
				"upgrade-to-" + targetOS,
				upgradeSteps[targetOS],
				testInstance,
			})
		}

		// Only 2012r2 supports retry-upgrade.
		if targetOS == versionWindows2012r2 {
			tcs = append(tcs, testCase{
				"retry-upgrade-to-" + targetOS,
				retryUpgradeSteps[targetOS],
				testInstanceWithStartupScript,
			})
		}
	}

	for _, tc := range tcs {
		u := newTestUpgrader()
		u.Instance = daisyutils.GetInstanceURI(testProject, testZone, tc.instanceName)

		err := u.validateAndDeriveParams()
		if err != nil {
			t.Errorf("[%v]: validateAndDeriveParams failed: %v", tc.testName, err)
			continue
		}

		_, err = u.generateWorkflowWithSteps("test", DefaultTimeout, tc.populateFunc)
		assert.NoError(t, err, "[test name: %v] Unexpected error.", tc.testName)
	}
}

func TestRunWorkflowWithSteps(t *testing.T) {
	type testCase struct {
		testName                  string
		populateFunc              func(*upgrader, *daisy.Workflow) error
		expectExitOnPopulateSteps bool
	}

	tcs := []testCase{
		{"populate without error", func(u *upgrader, w *daisy.Workflow) error {
			w.Steps = map[string]*daisy.Step{
				"step1": {},
			}
			return nil
		}, false},
		{"populate with error", func(u *upgrader, w *daisy.Workflow) error {
			w.Steps = map[string]*daisy.Step{
				"step1": {},
			}
			return daisy.Errf("some error")
		}, true},
	}

	for _, tc := range tcs {
		u := newTestUpgrader()
		err := u.validateAndDeriveParams()
		if err != nil {
			t.Errorf("[%v]: validateAndDeriveParams failed: %v", tc.testName, err)
			continue
		}

		worker, err := u.runWorkflowWithSteps("test", "10m", tc.populateFunc)

		daisyutils.CheckWorkflow(worker, func(w *daisy.Workflow, workflowCreationError error) {
			if _, ok := w.Steps["step1"]; !ok {
				t.Errorf("[%v]: missed step.", tc.testName)
			}

			if tc.expectExitOnPopulateSteps {
				if w.DefaultTimeout == u.Timeout {
					t.Errorf("[%v]: Default timeout of the workflow '%v' was overrided to '%v' unexpected.", tc.testName, w.DefaultTimeout, u.Timeout)
				}
			} else if w.DefaultTimeout != u.Timeout {
				t.Errorf("[%v]: Default timeout of the workflow '%v' should be overrided to '%v'.", tc.testName, w.DefaultTimeout, u.Timeout)
			}
		})
	}
}

func TestRunAllWorkflowFunctions(t *testing.T) {
	u := newTestUpgrader()

	type testCase struct {
		testName     string
		workflowFunc func() (daisyutils.DaisyWorker, error)
	}

	tcs := []testCase{
		{"prepare", u.prepare},
		{"upgrade", u.upgrade},
		{"reboot", u.reboot},
		{"retry-upgrade", u.retryUpgrade},
		{"rollback", u.rollback},
		{"cleanup", u.cleanup},
	}

	for _, tc := range tcs {
		err := u.validateAndDeriveParams()
		if err != nil {
			t.Errorf("[%v]: validateAndDeriveParams failed: %v", tc.testName, err)
			continue
		}

		worker, err := tc.workflowFunc()

		daisyutils.CheckWorkflow(worker, func(w *daisy.Workflow, workflowCreationError error) {
			if w.DefaultTimeout != u.Timeout {
				t.Errorf("[%v]: Default timeout of the workflow '%v' should be overrided to '%v'.", tc.testName, w.DefaultTimeout, u.Timeout)
			}
		})
	}
}