From da95bb3f720763928b5827d4a4d2b078b1096043 Mon Sep 17 00:00:00 2001
From: zackattack01 <zack@kolide.co>
Date: Mon, 23 Dec 2024 11:50:01 -0500
Subject: [PATCH] fix up agent reset_test.go defers

---
 ee/agent/reset_test.go | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/ee/agent/reset_test.go b/ee/agent/reset_test.go
index 4514f06f2..96595bee6 100644
--- a/ee/agent/reset_test.go
+++ b/ee/agent/reset_test.go
@@ -35,12 +35,12 @@ func TestMain(m *testing.M) {
 		fmt.Printf("failed to make temp dir for test osquery binary: %v", err)
 		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit inside tests
 	}
-	defer os.RemoveAll(downloadDir)
 
 	target := packaging.Target{}
 	if err := target.PlatformFromString(runtime.GOOS); err != nil {
 		fmt.Printf("error parsing platform %s: %v", runtime.GOOS, err)
-		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit inside tests
+		os.RemoveAll(downloadDir) // explicit removal as defer will not run when os.Exit is called
+		os.Exit(1)                //nolint:forbidigo // Fine to use os.Exit inside tests
 	}
 	target.Arch = packaging.ArchFlavor(runtime.GOARCH)
 	if runtime.GOOS == "darwin" {
@@ -48,12 +48,13 @@ func TestMain(m *testing.M) {
 	}
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
 
 	dlPath, err := packaging.FetchBinary(ctx, downloadDir, "osqueryd", target.PlatformBinaryName("osqueryd"), "nightly", target)
 	if err != nil {
 		fmt.Printf("error fetching binary osqueryd binary: %v", err)
-		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit inside tests
+		cancel()                  // explicit cancel as defer will not run when os.Exit is called
+		os.RemoveAll(downloadDir) // explicit removal as defer will not run when os.Exit is called
+		os.Exit(1)                //nolint:forbidigo // Fine to use os.Exit inside tests
 	}
 
 	testOsqueryBinary = filepath.Join(downloadDir, filepath.Base(dlPath))
@@ -63,12 +64,17 @@ func TestMain(m *testing.M) {
 
 	if err := fsutil.CopyFile(dlPath, testOsqueryBinary); err != nil {
 		fmt.Printf("error copying osqueryd binary: %v", err)
-		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit inside tests
+		cancel()                  // explicit cancel as defer will not run when os.Exit is called
+		os.RemoveAll(downloadDir) // explicit removal as defer will not run when os.Exit is called
+		os.Exit(1)                //nolint:forbidigo // Fine to use os.Exit inside tests
 	}
 
 	// Run the tests
 	retCode := m.Run()
-	os.Exit(retCode) //nolint:forbidigo // Fine to use os.Exit inside tests
+
+	cancel()                  // explicit cancel as defer will not run when os.Exit is called
+	os.RemoveAll(downloadDir) // explicit removal as defer will not run when os.Exit is called
+	os.Exit(retCode)          //nolint:forbidigo // Fine to use os.Exit inside tests
 }
 
 func TestDetectAndRemediateHardwareChange(t *testing.T) {