runtime and agent test improvements (kolide#2014)

RebeccaMahany · Dec 23, 2024 · a5c83f1 · a5c83f1
1 parent 730f144
commit a5c83f1
Show file tree

Hide file tree

Showing 5 changed files with 41 additions and 28 deletions.
diff --git a/ee/agent/reset_test.go b/ee/agent/reset_test.go
@@ -35,25 +35,26 @@ 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" {
 		target.Arch = packaging.Universal
 	}
 
 	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) {

diff --git a/pkg/osquery/runtime/osqueryinstance_test.go b/pkg/osquery/runtime/osqueryinstance_test.go
@@ -50,7 +50,7 @@ func TestCreateOsqueryCommand(t *testing.T) {
 		extensionAutoloadPath: "/foo/bar/osquery.autoload",
 	}
 
-	osquerydPath := testOsqueryBinaryDirectory
+	osquerydPath := testOsqueryBinary
 
 	k := typesMocks.NewKnapsack(t)
 	k.On("WatchdogEnabled").Return(true)
@@ -86,7 +86,7 @@ func TestCreateOsqueryCommandWithFlags(t *testing.T) {
 	i := newInstance(types.DefaultRegistrationID, k, mockServiceClient())
 
 	cmd, err := i.createOsquerydCommand(
-		testOsqueryBinaryDirectory,
+		testOsqueryBinary,
 		&osqueryFilePaths{},
 	)
 	require.NoError(t, err)
@@ -120,7 +120,7 @@ func TestCreateOsqueryCommand_SetsEnabledWatchdogSettingsAppropriately(t *testin
 	i := newInstance(types.DefaultRegistrationID, k, mockServiceClient())
 
 	cmd, err := i.createOsquerydCommand(
-		testOsqueryBinaryDirectory,
+		testOsqueryBinary,
 		&osqueryFilePaths{},
 	)
 	require.NoError(t, err)
@@ -170,7 +170,7 @@ func TestCreateOsqueryCommand_SetsDisabledWatchdogSettingsAppropriately(t *testi
 	i := newInstance(types.DefaultRegistrationID, k, mockServiceClient())
 
 	cmd, err := i.createOsquerydCommand(
-		testOsqueryBinaryDirectory,
+		testOsqueryBinary,
 		&osqueryFilePaths{},
 	)
 	require.NoError(t, err)

diff --git a/pkg/osquery/runtime/osqueryinstance_windows_test.go b/pkg/osquery/runtime/osqueryinstance_windows_test.go
@@ -16,7 +16,7 @@ import (
 func TestCreateOsqueryCommandEnvVars(t *testing.T) {
 	t.Parallel()
 
-	osquerydPath := testOsqueryBinaryDirectory
+	osquerydPath := testOsqueryBinary
 
 	k := typesMocks.NewKnapsack(t)
 	k.On("WatchdogEnabled").Return(true)

diff --git a/pkg/osquery/runtime/runtime_posix_test.go b/pkg/osquery/runtime/runtime_posix_test.go
@@ -50,7 +50,7 @@ func TestOsquerySlowStart(t *testing.T) {
 	k.On("OsqueryFlags").Return([]string{}).Maybe()
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("LoggingInterval").Return(5 * time.Minute).Maybe()
 	k.On("LogMaxBytesPerBatch").Return(0).Maybe()
 	k.On("Transport").Return("jsonrpc").Maybe()
@@ -97,7 +97,7 @@ func TestExtensionSocketPath(t *testing.T) {
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryVerbose").Return(true).Maybe()
 	k.On("OsqueryFlags").Return([]string{}).Maybe()
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("LoggingInterval").Return(5 * time.Minute).Maybe()
 	k.On("LogMaxBytesPerBatch").Return(0).Maybe()
 	k.On("Transport").Return("jsonrpc").Maybe()

diff --git a/pkg/osquery/runtime/runtime_test.go b/pkg/osquery/runtime/runtime_test.go
@@ -40,7 +40,7 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-var testOsqueryBinaryDirectory string
+var testOsqueryBinary string
 
 // TestMain overrides the default test main function. This allows us to share setup/teardown.
 func TestMain(m *testing.M) {
@@ -54,30 +54,37 @@ func TestMain(m *testing.M) {
 		fmt.Println("Failed to make temp dir for test binaries")
 		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit in tests
 	}
-	defer os.Remove(binDirectory)
 
 	s, err := storageci.NewStore(nil, multislogger.NewNopLogger(), storage.OsqueryHistoryInstanceStore.String())
 	if err != nil {
 		fmt.Println("Failed to make new store")
-		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit in tests
+		os.Remove(binDirectory) // explicit removal as defer will not run when os.Exit is called
+		os.Exit(1)              //nolint:forbidigo // Fine to use os.Exit in tests
 	}
 	if err := history.InitHistory(s); err != nil {
 		fmt.Println("Failed to init history")
-		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit in tests
+		os.Remove(binDirectory) // explicit removal as defer will not run when os.Exit is called
+		os.Exit(1)              //nolint:forbidigo // Fine to use os.Exit in tests
 	}
 
-	testOsqueryBinaryDirectory = filepath.Join(binDirectory, "osqueryd")
+	testOsqueryBinary = filepath.Join(binDirectory, "osqueryd")
+	if runtime.GOOS == "windows" {
+		testOsqueryBinary += ".exe"
+	}
 
 	thrift.ServerConnectivityCheckInterval = 100 * time.Millisecond
 
 	if err := downloadOsqueryInBinDir(binDirectory); err != nil {
 		fmt.Printf("Failed to download osquery: %v\n", err)
-		os.Exit(1) //nolint:forbidigo // Fine to use os.Exit in tests
+		os.Remove(binDirectory) // explicit removal as defer will not run when os.Exit is called
+		os.Exit(1)              //nolint:forbidigo // Fine to use os.Exit in tests
 	}
 
 	// Run the tests!
 	retCode := m.Run()
-	os.Exit(retCode) //nolint:forbidigo // Fine to use os.Exit in tests
+
+	os.Remove(binDirectory) // explicit removal as defer will not run when os.Exit is called
+	os.Exit(retCode)        //nolint:forbidigo // Fine to use os.Exit in tests
 }
 
 // downloadOsqueryInBinDir downloads osqueryd. This allows the test
@@ -163,7 +170,7 @@ func TestWithOsqueryFlags(t *testing.T) {
 	k.On("WatchdogEnabled").Return(false)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryFlags").Return([]string{"verbose=false"})
 	k.On("OsqueryVerbose").Return(false)
@@ -197,7 +204,7 @@ func TestFlagsChanged(t *testing.T) {
 	k.On("WatchdogDelaySec").Return(120)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryFlags").Return([]string{"verbose=false"})
 	k.On("OsqueryVerbose").Return(false)
@@ -328,7 +335,7 @@ func TestSimplePath(t *testing.T) {
 	k.On("WatchdogEnabled").Return(false)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryFlags").Return([]string{})
 	k.On("OsqueryVerbose").Return(true)
@@ -364,7 +371,7 @@ func TestMultipleInstances(t *testing.T) {
 	k.On("WatchdogEnabled").Return(false)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryFlags").Return([]string{})
 	k.On("OsqueryVerbose").Return(true)
@@ -423,7 +430,7 @@ func TestRunnerHandlesImmediateShutdownWithMultipleInstances(t *testing.T) {
 	k.On("WatchdogEnabled").Return(false)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryFlags").Return([]string{})
 	k.On("OsqueryVerbose").Return(true)
@@ -474,7 +481,7 @@ func TestMultipleShutdowns(t *testing.T) {
 	k.On("WatchdogEnabled").Return(false)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryFlags").Return([]string{})
 	k.On("OsqueryVerbose").Return(true)
@@ -506,7 +513,7 @@ func TestOsqueryDies(t *testing.T) {
 	k.On("WatchdogEnabled").Return(false)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything)
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory)
 	k.On("OsqueryFlags").Return([]string{})
 	k.On("OsqueryVerbose").Return(true)
@@ -608,7 +615,7 @@ func setupOsqueryInstanceForTests(t *testing.T) (runner *Runner, logBytes *threa
 	k.On("WatchdogDelaySec").Return(120)
 	k.On("RegisterChangeObserver", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Maybe()
 	k.On("Slogger").Return(slogger)
-	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinaryDirectory)
+	k.On("LatestOsquerydPath", mock.Anything).Return(testOsqueryBinary)
 	k.On("RootDirectory").Return(rootDirectory).Maybe()
 	k.On("OsqueryFlags").Return([]string{}).Maybe()
 	k.On("OsqueryVerbose").Return(true).Maybe()