From 1e675cabb5660d8ec6c5dc52f9425af8c93aebba Mon Sep 17 00:00:00 2001
From: henrygd <hank@henrygd.me>
Date: Fri, 25 Jul 2025 13:37:23 -0400
Subject: [PATCH] refactor agent data directory resolution (#991)

---
 beszel/internal/agent/data_dir.go      | 31 +++++++++++---------------
 beszel/internal/agent/data_dir_test.go | 10 ++++-----
 2 files changed, 18 insertions(+), 23 deletions(-)

diff --git a/beszel/internal/agent/data_dir.go b/beszel/internal/agent/data_dir.go
index 735429b..cb71366 100644
--- a/beszel/internal/agent/data_dir.go
+++ b/beszel/internal/agent/data_dir.go
@@ -9,35 +9,30 @@ import (
 )
 
 // getDataDir returns the path to the data directory for the agent and an error
-// if the directory is not valid. Pass an empty string to attempt to find the
-// optimal data directory.
-func getDataDir(dataDir string) (string, error) {
-	if dataDir == "" {
-		dataDir, _ = GetEnv("DATA_DIR")
+// if the directory is not valid. Attempts to find the optimal data directory if
+// no data directories are provided.
+func getDataDir(dataDirs ...string) (string, error) {
+	if len(dataDirs) > 0 {
+		return testDataDirs(dataDirs)
 	}
 
+	dataDir, _ := GetEnv("DATA_DIR")
 	if dataDir != "" {
-		return testDataDirs([]string{dataDir})
+		dataDirs = append(dataDirs, dataDir)
 	}
 
-	var dirsToTry []string
-
 	if runtime.GOOS == "windows" {
-		dirsToTry = []string{
+		dataDirs = append(dataDirs,
 			filepath.Join(os.Getenv("APPDATA"), "beszel-agent"),
 			filepath.Join(os.Getenv("LOCALAPPDATA"), "beszel-agent"),
-		}
+		)
 	} else {
-		homeDir, err := os.UserHomeDir()
-		if err != nil {
-			return "", err
-		}
-		dirsToTry = []string{
-			"/var/lib/beszel-agent",
-			filepath.Join(homeDir, ".config", "beszel"),
+		dataDirs = append(dataDirs, "/var/lib/beszel-agent")
+		if homeDir, err := os.UserHomeDir(); err == nil {
+			dataDirs = append(dataDirs, filepath.Join(homeDir, ".config", "beszel"))
 		}
 	}
-	return testDataDirs(dirsToTry)
+	return testDataDirs(dataDirs)
 }
 
 func testDataDirs(paths []string) (string, error) {
diff --git a/beszel/internal/agent/data_dir_test.go b/beszel/internal/agent/data_dir_test.go
index 74da801..7c12ffa 100644
--- a/beszel/internal/agent/data_dir_test.go
+++ b/beszel/internal/agent/data_dir_test.go
@@ -44,15 +44,15 @@ func TestGetDataDir(t *testing.T) {
 		oldValue := os.Getenv("DATA_DIR")
 		defer func() {
 			if oldValue == "" {
-				os.Unsetenv("DATA_DIR")
+				os.Unsetenv("BESZEL_AGENT_DATA_DIR")
 			} else {
-				os.Setenv("DATA_DIR", oldValue)
+				os.Setenv("BESZEL_AGENT_DATA_DIR", oldValue)
 			}
 		}()
 
-		os.Setenv("DATA_DIR", tempDir)
+		os.Setenv("BESZEL_AGENT_DATA_DIR", tempDir)
 
-		result, err := getDataDir("")
+		result, err := getDataDir()
 		require.NoError(t, err)
 		assert.Equal(t, tempDir, result)
 	})
@@ -79,7 +79,7 @@ func TestGetDataDir(t *testing.T) {
 
 		// This will try platform-specific defaults, which may or may not work
 		// We're mainly testing that it doesn't panic and returns some result
-		result, err := getDataDir("")
+		result, err := getDataDir()
 		// We don't assert success/failure here since it depends on system permissions
 		// Just verify we get a string result if no error
 		if err == nil {