From 29782756594b1c6045be6f773265110a792454d9 Mon Sep 17 00:00:00 2001
From: Subhransu <47723536+Shubhranshu153@users.noreply.github.com>
Date: Mon, 18 Nov 2024 12:00:47 -0800
Subject: [PATCH] fix: Refactor VM lifecycle test for flakiness (#1187)

Signed-off-by: Shubharanshu Mahapatra <shubhum@amazon.com>
---
 e2e/vm/lifecycle_test.go | 53 +++++++++++++++++++++++++++++++++-------
 1 file changed, 44 insertions(+), 9 deletions(-)

diff --git a/e2e/vm/lifecycle_test.go b/e2e/vm/lifecycle_test.go
index 940821b2b..fa34b9d38 100644
--- a/e2e/vm/lifecycle_test.go
+++ b/e2e/vm/lifecycle_test.go
@@ -6,16 +6,39 @@
 package vm
 
 import (
+	"runtime"
+	"time"
+
 	"github.com/onsi/ginkgo/v2"
 	"github.com/onsi/gomega"
 	"github.com/runfinch/common-tests/command"
 	"github.com/runfinch/common-tests/option"
 )
 
+var restoreVM = func(o *option.Option) {
+	origFinchCfg := readFile(finchConfigFilePath)
+	writeFile(finchConfigFilePath, origFinchCfg)
+	command.New(o, virtualMachineRootCmd, "stop", "-f").WithoutCheckingExitCode().WithTimeoutInSeconds(20).Run()
+	time.Sleep(1 * time.Second)
+	command.New(o, virtualMachineRootCmd, "remove", "-f").WithoutCheckingExitCode().WithTimeoutInSeconds(10).Run()
+	time.Sleep(1 * time.Second)
+	if runtime.GOOS == "windows" {
+		// clean up iptables
+		//nolint:lll // link to explanation
+		// https://docs.rancherdesktop.io/troubleshooting-tips/#q-how-do-i-fix-fata0005-subnet-1040024-overlaps-with-other-one-on-this-address-space-when-running-a-container-using-nerdctl-run
+		gomega.Expect(shutdownWSL()).Should(gomega.BeNil())
+	}
+	time.Sleep(1 * time.Second)
+	command.New(o, virtualMachineRootCmd, "init").WithoutCheckingExitCode().WithTimeoutInSeconds(160).Run()
+}
+
 var testVMLifecycle = func(o *option.Option) {
-	// These tests are run in serial because we only define one virtual machine instance.
 	ginkgo.Describe("virtual machine lifecycle", ginkgo.Serial, func() {
-		ginkgo.When("the virtual machine is in running status", func() {
+		ginkgo.Describe("when the virtual machine is in running status", func() {
+			ginkgo.AfterEach(func() {
+				restoreVM(o)
+			})
+
 			ginkgo.It("should fail to init/remove the virtual machine", func() {
 				command.New(o, virtualMachineRootCmd, "init").WithoutSuccessfulExit().Run()
 				command.New(o, virtualMachineRootCmd, "remove").WithoutSuccessfulExit().Run()
@@ -36,7 +59,6 @@ var testVMLifecycle = func(o *option.Option) {
 				command.Run(o, "images")
 				command.New(o, virtualMachineRootCmd, "remove", "--force").WithTimeoutInSeconds(160).Run()
 				command.RunWithoutSuccessfulExit(o, "images")
-				command.New(o, virtualMachineRootCmd, "init").WithTimeoutInSeconds(160).Run()
 			})
 
 			ginkgo.It("should be able to stop the virtual machine", func() {
@@ -46,7 +68,15 @@ var testVMLifecycle = func(o *option.Option) {
 			})
 		})
 
-		ginkgo.When("the virtual machine is in stopped status", func() {
+		ginkgo.Describe("when the virtual machine is in stopped status", func() {
+			ginkgo.BeforeEach(func() {
+				command.New(o, virtualMachineRootCmd, "stop", "--force").WithTimeoutInSeconds(90).Run()
+			})
+
+			ginkgo.AfterEach(func() {
+				restoreVM(o)
+			})
+
 			ginkgo.It("should fail to init/stop", func() {
 				command.New(o, virtualMachineRootCmd, "stop").WithoutSuccessfulExit().Run()
 				command.New(o, virtualMachineRootCmd, "init").WithoutSuccessfulExit().Run()
@@ -63,20 +93,25 @@ var testVMLifecycle = func(o *option.Option) {
 			})
 
 			ginkgo.It("should be able to remove the virtual machine", func() {
-				// don't asssume the VM will be in a stopped state (e.g. if the previous test fails)
-				command.New(o, virtualMachineRootCmd, "stop", "--force").WithTimeoutInSeconds(90).Run()
 				command.New(o, virtualMachineRootCmd, "remove").WithTimeoutInSeconds(160).Run()
 				command.RunWithoutSuccessfulExit(o, "images")
-				command.New(o, virtualMachineRootCmd, "init").WithTimeoutInSeconds(160).Run()
-				command.New(o, virtualMachineRootCmd, "stop").WithTimeoutInSeconds(90).Run()
 			})
+
 			ginkgo.It("should be able to force remove the virtual machine", func() {
 				command.New(o, virtualMachineRootCmd, "remove", "--force").WithTimeoutInSeconds(160).Run()
 				command.RunWithoutSuccessfulExit(o, "images")
 			})
 		})
 
-		ginkgo.When("the virtual machine instance does not exist", func() {
+		ginkgo.Describe("when the virtual machine instance does not exist", func() {
+			ginkgo.BeforeEach(func() {
+				command.New(o, virtualMachineRootCmd, "remove", "--force").WithTimeoutInSeconds(160).Run()
+			})
+
+			ginkgo.AfterEach(func() {
+				restoreVM(o)
+			})
+
 			ginkgo.It("should fail to start/stop", func() {
 				command.New(o, virtualMachineRootCmd, "start").WithoutSuccessfulExit().Run()
 				command.New(o, virtualMachineRootCmd, "stop").WithoutSuccessfulExit().Run()