From be70840609dd74ed5122bf7dbb46c90208c8ab8c Mon Sep 17 00:00:00 2001 From: henrygd Date: Fri, 20 Mar 2026 14:59:41 -0400 Subject: [PATCH] test: update tests that use os.Setenv to t.Setenv --- agent/client_test.go | 111 +++++++---------------------- agent/connection_manager_test.go | 42 ++++------- agent/data_dir_test.go | 23 +----- agent/disk_test.go | 14 +--- agent/gpu_test.go | 43 +++-------- agent/sensors_test.go | 31 +------- agent/server_test.go | 3 +- agent/smart_test.go | 2 +- agent/systemd_test.go | 14 +--- agent/utils/utils_test.go | 13 +--- internal/hub/agent_connect_test.go | 24 ++----- internal/hub/hub_test.go | 12 ++-- 12 files changed, 71 insertions(+), 261 deletions(-) diff --git a/agent/client_test.go b/agent/client_test.go index 09d6fb4f..931e16d0 100644 --- a/agent/client_test.go +++ b/agent/client_test.go @@ -70,19 +70,11 @@ func TestNewWebSocketClient(t *testing.T) { t.Run(tc.name, func(t *testing.T) { // Set up environment if tc.hubURL != "" { - os.Setenv("BESZEL_AGENT_HUB_URL", tc.hubURL) - } else { - os.Unsetenv("BESZEL_AGENT_HUB_URL") + t.Setenv("BESZEL_AGENT_HUB_URL", tc.hubURL) } if tc.token != "" { - os.Setenv("BESZEL_AGENT_TOKEN", tc.token) - } else { - os.Unsetenv("BESZEL_AGENT_TOKEN") + t.Setenv("BESZEL_AGENT_TOKEN", tc.token) } - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() client, err := newWebSocketClient(agent) @@ -138,12 +130,8 @@ func TestWebSocketClient_GetOptions(t *testing.T) { for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { // Set up environment - os.Setenv("BESZEL_AGENT_HUB_URL", tc.inputURL) - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", tc.inputURL) + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") client, err := newWebSocketClient(agent) require.NoError(t, err) @@ -185,12 +173,8 @@ func TestWebSocketClient_VerifySignature(t *testing.T) { require.NoError(t, err) // Set up environment - os.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") client, err := newWebSocketClient(agent) require.NoError(t, err) @@ -258,12 +242,8 @@ func TestWebSocketClient_HandleHubRequest(t *testing.T) { agent := createTestAgent(t) // Set up environment - os.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") client, err := newWebSocketClient(agent) require.NoError(t, err) @@ -350,13 +330,8 @@ func TestGetUserAgent(t *testing.T) { func TestWebSocketClient_Close(t *testing.T) { agent := createTestAgent(t) - // Set up environment - os.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") client, err := newWebSocketClient(agent) require.NoError(t, err) @@ -371,13 +346,8 @@ func TestWebSocketClient_Close(t *testing.T) { func TestWebSocketClient_ConnectRateLimit(t *testing.T) { agent := createTestAgent(t) - // Set up environment - os.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") client, err := newWebSocketClient(agent) require.NoError(t, err) @@ -393,20 +363,10 @@ func TestWebSocketClient_ConnectRateLimit(t *testing.T) { // TestGetToken tests the getToken function with various scenarios func TestGetToken(t *testing.T) { - unsetEnvVars := func() { - os.Unsetenv("BESZEL_AGENT_TOKEN") - os.Unsetenv("TOKEN") - os.Unsetenv("BESZEL_AGENT_TOKEN_FILE") - os.Unsetenv("TOKEN_FILE") - } - t.Run("token from TOKEN environment variable", func(t *testing.T) { - unsetEnvVars() - // Set TOKEN env var expectedToken := "test-token-from-env" - os.Setenv("TOKEN", expectedToken) - defer os.Unsetenv("TOKEN") + t.Setenv("TOKEN", expectedToken) token, err := getToken() assert.NoError(t, err) @@ -414,12 +374,9 @@ func TestGetToken(t *testing.T) { }) t.Run("token from BESZEL_AGENT_TOKEN environment variable", func(t *testing.T) { - unsetEnvVars() - // Set BESZEL_AGENT_TOKEN env var (should take precedence) expectedToken := "test-token-from-beszel-env" - os.Setenv("BESZEL_AGENT_TOKEN", expectedToken) - defer os.Unsetenv("BESZEL_AGENT_TOKEN") + t.Setenv("BESZEL_AGENT_TOKEN", expectedToken) token, err := getToken() assert.NoError(t, err) @@ -427,8 +384,6 @@ func TestGetToken(t *testing.T) { }) t.Run("token from TOKEN_FILE", func(t *testing.T) { - unsetEnvVars() - // Create a temporary token file expectedToken := "test-token-from-file" tokenFile, err := os.CreateTemp("", "token-test-*.txt") @@ -440,8 +395,7 @@ func TestGetToken(t *testing.T) { tokenFile.Close() // Set TOKEN_FILE env var - os.Setenv("TOKEN_FILE", tokenFile.Name()) - defer os.Unsetenv("TOKEN_FILE") + t.Setenv("TOKEN_FILE", tokenFile.Name()) token, err := getToken() assert.NoError(t, err) @@ -449,8 +403,6 @@ func TestGetToken(t *testing.T) { }) t.Run("token from BESZEL_AGENT_TOKEN_FILE", func(t *testing.T) { - unsetEnvVars() - // Create a temporary token file expectedToken := "test-token-from-beszel-file" tokenFile, err := os.CreateTemp("", "token-test-*.txt") @@ -462,8 +414,7 @@ func TestGetToken(t *testing.T) { tokenFile.Close() // Set BESZEL_AGENT_TOKEN_FILE env var (should take precedence) - os.Setenv("BESZEL_AGENT_TOKEN_FILE", tokenFile.Name()) - defer os.Unsetenv("BESZEL_AGENT_TOKEN_FILE") + t.Setenv("BESZEL_AGENT_TOKEN_FILE", tokenFile.Name()) token, err := getToken() assert.NoError(t, err) @@ -471,8 +422,6 @@ func TestGetToken(t *testing.T) { }) t.Run("TOKEN takes precedence over TOKEN_FILE", func(t *testing.T) { - unsetEnvVars() - // Create a temporary token file fileToken := "token-from-file" tokenFile, err := os.CreateTemp("", "token-test-*.txt") @@ -485,12 +434,8 @@ func TestGetToken(t *testing.T) { // Set both TOKEN and TOKEN_FILE envToken := "token-from-env" - os.Setenv("TOKEN", envToken) - os.Setenv("TOKEN_FILE", tokenFile.Name()) - defer func() { - os.Unsetenv("TOKEN") - os.Unsetenv("TOKEN_FILE") - }() + t.Setenv("TOKEN", envToken) + t.Setenv("TOKEN_FILE", tokenFile.Name()) token, err := getToken() assert.NoError(t, err) @@ -498,7 +443,10 @@ func TestGetToken(t *testing.T) { }) t.Run("error when neither TOKEN nor TOKEN_FILE is set", func(t *testing.T) { - unsetEnvVars() + t.Setenv("BESZEL_AGENT_TOKEN", "") + t.Setenv("TOKEN", "") + t.Setenv("BESZEL_AGENT_TOKEN_FILE", "") + t.Setenv("TOKEN_FILE", "") token, err := getToken() assert.Error(t, err) @@ -507,11 +455,8 @@ func TestGetToken(t *testing.T) { }) t.Run("error when TOKEN_FILE points to non-existent file", func(t *testing.T) { - unsetEnvVars() - // Set TOKEN_FILE to a non-existent file - os.Setenv("TOKEN_FILE", "/non/existent/file.txt") - defer os.Unsetenv("TOKEN_FILE") + t.Setenv("TOKEN_FILE", "/non/existent/file.txt") token, err := getToken() assert.Error(t, err) @@ -520,8 +465,6 @@ func TestGetToken(t *testing.T) { }) t.Run("handles empty token file", func(t *testing.T) { - unsetEnvVars() - // Create an empty token file tokenFile, err := os.CreateTemp("", "token-test-*.txt") require.NoError(t, err) @@ -529,8 +472,7 @@ func TestGetToken(t *testing.T) { tokenFile.Close() // Set TOKEN_FILE env var - os.Setenv("TOKEN_FILE", tokenFile.Name()) - defer os.Unsetenv("TOKEN_FILE") + t.Setenv("TOKEN_FILE", tokenFile.Name()) token, err := getToken() assert.NoError(t, err) @@ -538,8 +480,6 @@ func TestGetToken(t *testing.T) { }) t.Run("strips whitespace from TOKEN_FILE", func(t *testing.T) { - unsetEnvVars() - tokenWithWhitespace := " test-token-with-whitespace \n\t" expectedToken := "test-token-with-whitespace" tokenFile, err := os.CreateTemp("", "token-test-*.txt") @@ -550,8 +490,7 @@ func TestGetToken(t *testing.T) { require.NoError(t, err) tokenFile.Close() - os.Setenv("TOKEN_FILE", tokenFile.Name()) - defer os.Unsetenv("TOKEN_FILE") + t.Setenv("TOKEN_FILE", tokenFile.Name()) token, err := getToken() assert.NoError(t, err) diff --git a/agent/connection_manager_test.go b/agent/connection_manager_test.go index 8a848772..7a9ca4a0 100644 --- a/agent/connection_manager_test.go +++ b/agent/connection_manager_test.go @@ -7,7 +7,6 @@ import ( "fmt" "net" "net/url" - "os" "testing" "time" @@ -183,10 +182,6 @@ func TestConnectionManager_TickerManagement(t *testing.T) { // TestConnectionManager_WebSocketConnectionFlow tests WebSocket connection logic func TestConnectionManager_WebSocketConnectionFlow(t *testing.T) { - if testing.Short() { - t.Skip("Skipping WebSocket connection test in short mode") - } - agent := createTestAgent(t) cm := agent.connectionManager @@ -196,19 +191,18 @@ func TestConnectionManager_WebSocketConnectionFlow(t *testing.T) { assert.Equal(t, Disconnected, cm.State, "State should remain Disconnected after failed connection") // Test with invalid URL - os.Setenv("BESZEL_AGENT_HUB_URL", "invalid-url") - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() - - // Test with missing token - os.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") - os.Unsetenv("BESZEL_AGENT_TOKEN") + t.Setenv("BESZEL_AGENT_HUB_URL", "1,33%") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") _, err2 := newWebSocketClient(agent) - assert.Error(t, err2, "WebSocket client creation should fail without token") + assert.Error(t, err2, "WebSocket client creation should fail with invalid URL") + + // Test with missing token + t.Setenv("BESZEL_AGENT_HUB_URL", "http://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "") + + _, err3 := newWebSocketClient(agent) + assert.Error(t, err3, "WebSocket client creation should fail without token") } // TestConnectionManager_ReconnectionLogic tests reconnection prevention logic @@ -234,12 +228,8 @@ func TestConnectionManager_ConnectWithRateLimit(t *testing.T) { cm := agent.connectionManager // Set up environment for WebSocket client creation - os.Setenv("BESZEL_AGENT_HUB_URL", "ws://localhost:8080") - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", "ws://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") // Create WebSocket client wsClient, err := newWebSocketClient(agent) @@ -285,12 +275,8 @@ func TestConnectionManager_CloseWebSocket(t *testing.T) { }, "Should not panic when closing nil WebSocket client") // Set up environment and create WebSocket client - os.Setenv("BESZEL_AGENT_HUB_URL", "ws://localhost:8080") - os.Setenv("BESZEL_AGENT_TOKEN", "test-token") - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", "ws://localhost:8080") + t.Setenv("BESZEL_AGENT_TOKEN", "test-token") wsClient, err := newWebSocketClient(agent) require.NoError(t, err) diff --git a/agent/data_dir_test.go b/agent/data_dir_test.go index 66f97b8d..1705aacc 100644 --- a/agent/data_dir_test.go +++ b/agent/data_dir_test.go @@ -39,17 +39,7 @@ func TestGetDataDir(t *testing.T) { t.Run("DATA_DIR environment variable", func(t *testing.T) { tempDir := t.TempDir() - // Set environment variable - oldValue := os.Getenv("DATA_DIR") - defer func() { - if oldValue == "" { - os.Unsetenv("BESZEL_AGENT_DATA_DIR") - } else { - os.Setenv("BESZEL_AGENT_DATA_DIR", oldValue) - } - }() - - os.Setenv("BESZEL_AGENT_DATA_DIR", tempDir) + t.Setenv("BESZEL_AGENT_DATA_DIR", tempDir) result, err := GetDataDir() require.NoError(t, err) @@ -65,17 +55,6 @@ func TestGetDataDir(t *testing.T) { // Test fallback behavior (empty dataDir, no env var) t.Run("fallback to default directories", func(t *testing.T) { - // Clear DATA_DIR environment variable - oldValue := os.Getenv("DATA_DIR") - defer func() { - if oldValue == "" { - os.Unsetenv("DATA_DIR") - } else { - os.Setenv("DATA_DIR", oldValue) - } - }() - os.Unsetenv("DATA_DIR") - // 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() diff --git a/agent/disk_test.go b/agent/disk_test.go index 8a16d357..0c3efdad 100644 --- a/agent/disk_test.go +++ b/agent/disk_test.go @@ -687,18 +687,8 @@ func TestIsDockerSpecialMountpoint(t *testing.T) { } func TestInitializeDiskInfoWithCustomNames(t *testing.T) { - // Set up environment variables - oldEnv := os.Getenv("EXTRA_FILESYSTEMS") - defer func() { - if oldEnv != "" { - os.Setenv("EXTRA_FILESYSTEMS", oldEnv) - } else { - os.Unsetenv("EXTRA_FILESYSTEMS") - } - }() - // Test with custom names - os.Setenv("EXTRA_FILESYSTEMS", "sda1__my-storage,/dev/sdb1__backup-drive,nvme0n1p2") + t.Setenv("EXTRA_FILESYSTEMS", "sda1__my-storage,/dev/sdb1__backup-drive,nvme0n1p2") // Mock disk partitions (we'll just test the parsing logic) // Since the actual disk operations are system-dependent, we'll focus on the parsing @@ -726,7 +716,7 @@ func TestInitializeDiskInfoWithCustomNames(t *testing.T) { for _, tc := range testCases { t.Run("env_"+tc.envValue, func(t *testing.T) { - os.Setenv("EXTRA_FILESYSTEMS", tc.envValue) + t.Setenv("EXTRA_FILESYSTEMS", tc.envValue) // Create mock partitions that would match our test cases partitions := []disk.PartitionStat{} diff --git a/agent/gpu_test.go b/agent/gpu_test.go index 4cf0c1a5..ade17d91 100644 --- a/agent/gpu_test.go +++ b/agent/gpu_test.go @@ -1083,8 +1083,6 @@ func TestCalculateGPUAverage(t *testing.T) { func TestGPUCapabilitiesAndLegacyPriority(t *testing.T) { // Save original PATH - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) hasAmdSysfs := (&GPUManager{}).hasAmdSysfs() tests := []struct { @@ -1178,7 +1176,7 @@ echo "[]"` { name: "no gpu tools available", setupCommands: func(_ string) error { - os.Setenv("PATH", "") + t.Setenv("PATH", "") return nil }, wantErr: true, @@ -1188,7 +1186,7 @@ echo "[]"` for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { tempDir := t.TempDir() - os.Setenv("PATH", tempDir) + t.Setenv("PATH", tempDir) if err := tt.setupCommands(tempDir); err != nil { t.Fatal(err) } @@ -1234,13 +1232,9 @@ echo "[]"` } func TestCollectorStartHelpers(t *testing.T) { - // Save original PATH - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) - // Set up temp dir with the commands dir := t.TempDir() - os.Setenv("PATH", dir) + t.Setenv("PATH", dir) tests := []struct { name string @@ -1370,11 +1364,8 @@ echo '[{"device_name":"NVIDIA Test GPU","temp":"52C","power_draw":"31W","gpu_uti } func TestNewGPUManagerPriorityNvtopFallback(t *testing.T) { - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) - dir := t.TempDir() - os.Setenv("PATH", dir) + t.Setenv("PATH", dir) t.Setenv("BESZEL_AGENT_GPU_COLLECTOR", "nvtop,nvidia-smi") nvtopPath := filepath.Join(dir, "nvtop") @@ -1399,11 +1390,8 @@ echo "0, NVIDIA Priority GPU, 45, 512, 2048, 12, 25"` } func TestNewGPUManagerPriorityMixedCollectors(t *testing.T) { - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) - dir := t.TempDir() - os.Setenv("PATH", dir) + t.Setenv("PATH", dir) t.Setenv("BESZEL_AGENT_GPU_COLLECTOR", "intel_gpu_top,rocm-smi") intelPath := filepath.Join(dir, "intel_gpu_top") @@ -1433,11 +1421,8 @@ echo '{"card0": {"Temperature (Sensor edge) (C)": "49.0", "Current Socket Graphi } func TestNewGPUManagerPriorityNvmlFallbackToNvidiaSmi(t *testing.T) { - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) - dir := t.TempDir() - os.Setenv("PATH", dir) + t.Setenv("PATH", dir) t.Setenv("BESZEL_AGENT_GPU_COLLECTOR", "nvml,nvidia-smi") nvidiaPath := filepath.Join(dir, "nvidia-smi") @@ -1456,11 +1441,8 @@ echo "0, NVIDIA Fallback GPU, 41, 256, 1024, 8, 14"` } func TestNewGPUManagerConfiguredCollectorsMustStart(t *testing.T) { - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) - dir := t.TempDir() - os.Setenv("PATH", dir) + t.Setenv("PATH", dir) t.Run("configured valid collector unavailable", func(t *testing.T) { t.Setenv("BESZEL_AGENT_GPU_COLLECTOR", "nvidia-smi") @@ -1480,11 +1462,8 @@ func TestNewGPUManagerConfiguredCollectorsMustStart(t *testing.T) { } func TestNewGPUManagerJetsonIgnoresCollectorConfig(t *testing.T) { - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) - dir := t.TempDir() - os.Setenv("PATH", dir) + t.Setenv("PATH", dir) t.Setenv("BESZEL_AGENT_GPU_COLLECTOR", "nvidia-smi") tegraPath := filepath.Join(dir, "tegrastats") @@ -1719,12 +1698,8 @@ func TestIntelUpdateFromStats(t *testing.T) { } func TestIntelCollectorStreaming(t *testing.T) { - // Save and override PATH - origPath := os.Getenv("PATH") - defer os.Setenv("PATH", origPath) - dir := t.TempDir() - os.Setenv("PATH", dir) + t.Setenv("PATH", dir) // Create a fake intel_gpu_top that prints -l format with four samples (first will be skipped) and exits scriptPath := filepath.Join(dir, "intel_gpu_top") diff --git a/agent/sensors_test.go b/agent/sensors_test.go index a055e5e6..de27ca35 100644 --- a/agent/sensors_test.go +++ b/agent/sensors_test.go @@ -5,7 +5,6 @@ package agent import ( "context" "fmt" - "os" "testing" "github.com/henrygd/beszel/internal/entities/system" @@ -329,34 +328,10 @@ func TestNewSensorConfigWithEnv(t *testing.T) { } func TestNewSensorConfig(t *testing.T) { - // Save original environment variables - originalPrimary, hasPrimary := os.LookupEnv("BESZEL_AGENT_PRIMARY_SENSOR") - originalSys, hasSys := os.LookupEnv("BESZEL_AGENT_SYS_SENSORS") - originalSensors, hasSensors := os.LookupEnv("BESZEL_AGENT_SENSORS") - - // Restore environment variables after the test - defer func() { - // Clean up test environment variables - os.Unsetenv("BESZEL_AGENT_PRIMARY_SENSOR") - os.Unsetenv("BESZEL_AGENT_SYS_SENSORS") - os.Unsetenv("BESZEL_AGENT_SENSORS") - - // Restore original values if they existed - if hasPrimary { - os.Setenv("BESZEL_AGENT_PRIMARY_SENSOR", originalPrimary) - } - if hasSys { - os.Setenv("BESZEL_AGENT_SYS_SENSORS", originalSys) - } - if hasSensors { - os.Setenv("BESZEL_AGENT_SENSORS", originalSensors) - } - }() - // Set test environment variables - os.Setenv("BESZEL_AGENT_PRIMARY_SENSOR", "test_primary") - os.Setenv("BESZEL_AGENT_SYS_SENSORS", "/test/path") - os.Setenv("BESZEL_AGENT_SENSORS", "test_sensor1,test_*,test_sensor3") + t.Setenv("BESZEL_AGENT_PRIMARY_SENSOR", "test_primary") + t.Setenv("BESZEL_AGENT_SYS_SENSORS", "/test/path") + t.Setenv("BESZEL_AGENT_SENSORS", "test_sensor1,test_*,test_sensor3") agent := &Agent{} result := agent.newSensorConfig() diff --git a/agent/server_test.go b/agent/server_test.go index acdeffa7..1cd52626 100644 --- a/agent/server_test.go +++ b/agent/server_test.go @@ -183,8 +183,7 @@ func TestStartServer(t *testing.T) { } func TestStartServerDisableSSH(t *testing.T) { - os.Setenv("BESZEL_AGENT_DISABLE_SSH", "true") - defer os.Unsetenv("BESZEL_AGENT_DISABLE_SSH") + t.Setenv("BESZEL_AGENT_DISABLE_SSH", "true") agent, err := NewAgent("") require.NoError(t, err) diff --git a/agent/smart_test.go b/agent/smart_test.go index 478e96e0..ec6eb5d9 100644 --- a/agent/smart_test.go +++ b/agent/smart_test.go @@ -1035,7 +1035,7 @@ func TestRefreshExcludedDevices(t *testing.T) { t.Setenv("EXCLUDE_SMART", tt.envValue) } else { // Ensure env var is not set for empty test - os.Unsetenv("EXCLUDE_SMART") + t.Setenv("EXCLUDE_SMART", "") } sm := &SmartManager{} diff --git a/agent/systemd_test.go b/agent/systemd_test.go index 685f9ac7..5810468b 100644 --- a/agent/systemd_test.go +++ b/agent/systemd_test.go @@ -167,16 +167,12 @@ func TestGetServicePatterns(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - // Clean up any existing env vars - os.Unsetenv("BESZEL_AGENT_SERVICE_PATTERNS") - os.Unsetenv("SERVICE_PATTERNS") - // Set up environment variables if tt.prefixedEnv != "" { - os.Setenv("BESZEL_AGENT_SERVICE_PATTERNS", tt.prefixedEnv) + t.Setenv("BESZEL_AGENT_SERVICE_PATTERNS", tt.prefixedEnv) } if tt.unprefixedEnv != "" { - os.Setenv("SERVICE_PATTERNS", tt.unprefixedEnv) + t.Setenv("SERVICE_PATTERNS", tt.unprefixedEnv) } // Run the function @@ -184,12 +180,6 @@ func TestGetServicePatterns(t *testing.T) { // Verify results assert.Equal(t, tt.expected, result, "Patterns should match expected values") - - // Cleanup - if tt.cleanupEnvVars { - os.Unsetenv("BESZEL_AGENT_SERVICE_PATTERNS") - os.Unsetenv("SERVICE_PATTERNS") - } }) } } diff --git a/agent/utils/utils_test.go b/agent/utils/utils_test.go index 9231a6c4..a7b51804 100644 --- a/agent/utils/utils_test.go +++ b/agent/utils/utils_test.go @@ -134,10 +134,8 @@ func TestGetEnv(t *testing.T) { prefixedKey := "BESZEL_AGENT_" + key t.Run("prefixed variable exists", func(t *testing.T) { - os.Setenv(prefixedKey, "prefixed_val") - os.Setenv(key, "unprefixed_val") - defer os.Unsetenv(prefixedKey) - defer os.Unsetenv(key) + t.Setenv(prefixedKey, "prefixed_val") + t.Setenv(key, "unprefixed_val") val, exists := GetEnv(key) assert.True(t, exists) @@ -145,9 +143,7 @@ func TestGetEnv(t *testing.T) { }) t.Run("only unprefixed variable exists", func(t *testing.T) { - os.Unsetenv(prefixedKey) - os.Setenv(key, "unprefixed_val") - defer os.Unsetenv(key) + t.Setenv(key, "unprefixed_val") val, exists := GetEnv(key) assert.True(t, exists) @@ -155,9 +151,6 @@ func TestGetEnv(t *testing.T) { }) t.Run("neither variable exists", func(t *testing.T) { - os.Unsetenv(prefixedKey) - os.Unsetenv(key) - val, exists := GetEnv(key) assert.False(t, exists) assert.Empty(t, val) diff --git a/internal/hub/agent_connect_test.go b/internal/hub/agent_connect_test.go index 2ec8d041..8ab9cada 100644 --- a/internal/hub/agent_connect_test.go +++ b/internal/hub/agent_connect_test.go @@ -898,12 +898,8 @@ func TestAgentWebSocketIntegration(t *testing.T) { require.NoError(t, err) // Set up environment variables for the agent - os.Setenv("BESZEL_AGENT_HUB_URL", ts.URL) - os.Setenv("BESZEL_AGENT_TOKEN", tc.agentToken) - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", ts.URL) + t.Setenv("BESZEL_AGENT_TOKEN", tc.agentToken) // Start agent in background done := make(chan error, 1) @@ -1081,12 +1077,8 @@ func TestMultipleSystemsWithSameUniversalToken(t *testing.T) { require.NoError(t, err) // Set up environment variables for the agent - os.Setenv("BESZEL_AGENT_HUB_URL", ts.URL) - os.Setenv("BESZEL_AGENT_TOKEN", universalToken) - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", ts.URL) + t.Setenv("BESZEL_AGENT_TOKEN", universalToken) // Count systems before connection systemsBefore, err := testApp.FindRecordsByFilter("systems", "users ~ {:userId}", "", -1, 0, map[string]any{"userId": userRecord.Id}) @@ -1244,12 +1236,8 @@ func TestPermanentUniversalTokenFromDB(t *testing.T) { require.NoError(t, err) // Set up environment variables for the agent - os.Setenv("BESZEL_AGENT_HUB_URL", ts.URL) - os.Setenv("BESZEL_AGENT_TOKEN", universalToken) - defer func() { - os.Unsetenv("BESZEL_AGENT_HUB_URL") - os.Unsetenv("BESZEL_AGENT_TOKEN") - }() + t.Setenv("BESZEL_AGENT_HUB_URL", ts.URL) + t.Setenv("BESZEL_AGENT_TOKEN", universalToken) // Start agent in background done := make(chan error, 1) diff --git a/internal/hub/hub_test.go b/internal/hub/hub_test.go index 40e35f17..a7a36e19 100644 --- a/internal/hub/hub_test.go +++ b/internal/hub/hub_test.go @@ -733,10 +733,8 @@ func TestFirstUserCreation(t *testing.T) { }) t.Run("CreateUserEndpoint not available when USER_EMAIL, USER_PASSWORD are set", func(t *testing.T) { - os.Setenv("BESZEL_HUB_USER_EMAIL", "me@example.com") - os.Setenv("BESZEL_HUB_USER_PASSWORD", "password123") - defer os.Unsetenv("BESZEL_HUB_USER_EMAIL") - defer os.Unsetenv("BESZEL_HUB_USER_PASSWORD") + t.Setenv("BESZEL_HUB_USER_EMAIL", "me@example.com") + t.Setenv("BESZEL_HUB_USER_PASSWORD", "password123") hub, _ := beszelTests.NewTestHub(t.TempDir()) defer hub.Cleanup() @@ -852,13 +850,12 @@ func TestAutoLoginMiddleware(t *testing.T) { var hubs []*beszelTests.TestHub defer func() { - defer os.Unsetenv("AUTO_LOGIN") for _, hub := range hubs { hub.Cleanup() } }() - os.Setenv("AUTO_LOGIN", "user@test.com") + t.Setenv("AUTO_LOGIN", "user@test.com") testAppFactory := func(t testing.TB) *pbTests.TestApp { hub, _ := beszelTests.NewTestHub(t.TempDir()) @@ -906,13 +903,12 @@ func TestTrustedHeaderMiddleware(t *testing.T) { var hubs []*beszelTests.TestHub defer func() { - defer os.Unsetenv("TRUSTED_AUTH_HEADER") for _, hub := range hubs { hub.Cleanup() } }() - os.Setenv("TRUSTED_AUTH_HEADER", "X-Beszel-Trusted") + t.Setenv("TRUSTED_AUTH_HEADER", "X-Beszel-Trusted") testAppFactory := func(t testing.TB) *pbTests.TestApp { hub, _ := beszelTests.NewTestHub(t.TempDir())