3 files changed, 44 insertions, 11 deletions
diff --git a/tests/e2e/lifecycle_test.go b/tests/e2e/lifecycle_test.go
index 6cff0c8..ffd731f 100644
--- a/tests/e2e/lifecycle_test.go
+++ b/tests/e2e/lifecycle_test.go
@@ -1,7 +1,6 @@
 package e2e
 
 import (
-	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -48,7 +47,9 @@ func waitForLifecycle(t *testing.T, binaryPath, runtimeDir, profile string, expe
 			t.Fatalf("Timed out waiting for lifecycle state: expected active=%v", expectedActive)
 		case <-tick.C:
 			cmd := exec.Command(binaryPath, "test-lifecycle", "--profile", profile)
-			cmd.Env = append(os.Environ(), fmt.Sprintf("XDG_RUNTIME_DIR=%s", runtimeDir))
+			cmd.Env = SetEnvOverrides(map[string]string{"XDG_RUNTIME_DIR": runtimeDir})
+			cmd.Stdout = os.Stdout
+			cmd.Stderr = os.Stderr
 			err := cmd.Run()
 
 			isActive := err == nil
@@ -79,7 +80,9 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 	t.Run("ReferenceCounting", func(t *testing.T) {
 		// Start a process that exits quickly
 		cmd1 := exec.Command(binaryPath, "--profile", "default", "--", "sleep", "1.0")
-		cmd1.Env = append(os.Environ(), fmt.Sprintf("XDG_RUNTIME_DIR=%s", tmpRuntimeDir))
+		cmd1.Env = SetEnvOverrides(map[string]string{"XDG_RUNTIME_DIR": tmpRuntimeDir})
+		cmd1.Stdout = os.Stdout
+		cmd1.Stderr = os.Stderr
 		if err := cmd1.Start(); err != nil {
 			t.Fatalf("Failed to start cmd1: %v", err)
 		}
@@ -87,9 +90,11 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 		// Verify PID file exists using polling
 		waitForLifecycle(t, binaryPath, tmpRuntimeDir, "default", true)
 
-		// Start a second process using the same profile
-		cmd2 := exec.Command(binaryPath, "--profile", "default", "--", "sleep", "1.0")
-		cmd2.Env = append(os.Environ(), fmt.Sprintf("XDG_RUNTIME_DIR=%s", tmpRuntimeDir))
+		// Start a second process using the same profile with a longer sleep
+		cmd2 := exec.Command(binaryPath, "--profile", "default", "--", "sleep", "5.0")
+		cmd2.Env = SetEnvOverrides(map[string]string{"XDG_RUNTIME_DIR": tmpRuntimeDir})
+		cmd2.Stdout = os.Stdout
+		cmd2.Stderr = os.Stderr
 		if err := cmd2.Start(); err != nil {
 			t.Fatalf("Failed to start cmd2: %v", err)
 		}
@@ -100,7 +105,7 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 			t.Fatalf("cmd1 failed: %v", err)
 		}
 
-		// Poll for the count to drop back to 1
+		// Poll for the count to drop back to 1 (cmd2 should still be running)
 		waitForLifecycle(t, binaryPath, tmpRuntimeDir, "default", true)
 
 		// Wait for second process to exit naturally
diff --git a/tests/e2e/mount_leak_test.go b/tests/e2e/mount_leak_test.go
index bdc9d75..428675f 100644
--- a/tests/e2e/mount_leak_test.go
+++ b/tests/e2e/mount_leak_test.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"os/exec"
+	"path/filepath"
 	"strings"
 	"testing"
 )
@@ -21,11 +22,11 @@ func TestDNSMountLeak(t *testing.T) {
 	dnsServer := "8.8.8.8"
 
 	// Pre-create a dummy config for the profile
-	configDir := "/tmp/wg-wrap-test-configs"
-	if err := os.MkdirAll(configDir, 0755); err != nil {
-		t.Fatalf("failed to create config dir: %v", err)
+	configDir := t.TempDir()
+	if err := os.MkdirAll(filepath.Join(configDir, "profiles"), 0755); err != nil {
+		t.Fatalf("failed to create profiles dir: %v", err)
 	}
-	configPath := fmt.Sprintf("%s/%s.conf", configDir, profile)
+	configPath := filepath.Join(configDir, "profiles", profile+".conf")
 	if err := os.WriteFile(configPath, []byte("[Interface]\nAddress = 10.0.0.1/24\nPrivateKey = aAAA\n"), 0644); err != nil {
 		t.Fatalf("failed to write config file: %v", err)
 	}
diff --git a/tests/e2e/test_helpers.go b/tests/e2e/test_helpers.go
index 0ae83aa..6d65011 100644
--- a/tests/e2e/test_helpers.go
+++ b/tests/e2e/test_helpers.go
@@ -3,6 +3,7 @@ package e2e
 import (
 	"fmt"
 	"os"
+	"strings"
 )
 
 // GetBinaryPath resolves the path to the wg-wrap binary.
@@ -19,3 +20,29 @@ func GetBinaryPath() (string, error) {
 
 	return path, nil
 }
+
+// SetEnvOverrides returns a new slice of environment variables with the provided overrides applied.
+// It ensures that overriding an existing variable replaces it rather than appending it.
+func SetEnvOverrides(overrides map[string]string) []string {
+	env := os.Environ()
+	newEnv := make([]string, 0, len(env)+len(overrides))
+
+	for _, e := range env {
+		matched := false
+		for k := range overrides {
+			if strings.HasPrefix(e, k+"=") {
+				matched = true
+				break
+			}
+		}
+		if !matched {
+			newEnv = append(newEnv, e)
+		}
+	}
+
+	for k, v := range overrides {
+		newEnv = append(newEnv, fmt.Sprintf("%s=%s", k, v))
+	}
+
+	return newEnv
+}