1 files changed, 68 insertions, 37 deletions

diff --git a/tests/e2e/lifecycle_test.go b/tests/e2e/lifecycle_test.go
index baf9f56..649dbc0 100644
--- a/tests/e2e/lifecycle_test.go
+++ b/tests/e2e/lifecycle_test.go
@@ -9,6 +9,34 @@ import (
 	"time"
 )
 
+func waitForPids(t *testing.T, pidsDir string, expectedCount int) {
+	timeout := time.After(2 * time.Second)
+	tick := time.NewTicker(50 * time.Millisecond)
+	defer tick.Stop()
+
+	for {
+		select {
+		case <-timeout:
+			files, err := os.ReadDir(pidsDir)
+			if err != nil {
+				t.Fatalf("Failed to read pids dir during timeout: %v", err)
+			}
+			t.Fatalf("Timed out waiting for %d PID files, got %d", expectedCount, len(files))
+		case <-tick.C:
+			files, err := os.ReadDir(pidsDir)
+			if err != nil {
+				if os.IsNotExist(err) {
+					continue
+				}
+				t.Fatalf("Failed to read pids dir: %v", err)
+			}
+			if len(files) == expectedCount {
+				return
+			}
+		}
+	}
+}
+
 func TestNamespaceLifecycleAutomation(t *testing.T) {
 	// 1. Setup Environment
 	binaryPath, err := GetBinaryPath()
@@ -32,17 +60,8 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 			t.Fatalf("Failed to start cmd1: %v", err)
 		}
 
-		// Allow a moment for the bootstrap loop to complete and register the PID
-		time.Sleep(500 * time.Millisecond)
-
-		// Verify PID file exists
-		files, err := os.ReadDir(pidsDir)
-		if err != nil {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if len(files) != 1 {
-			t.Errorf("Expected 1 PID file, got %d", len(files))
-		}
+		// Verify PID file exists using polling
+		waitForPids(t, pidsDir, 1)
 
 		// Start a second process using the same profile
 		cmd2 := exec.Command(binaryPath, "--profile", profile, "--", "sleep", "0.1")
@@ -50,43 +69,55 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 		if err := cmd2.Start(); err != nil {
 			t.Fatalf("Failed to start cmd2: %v", err)
 		}
-		time.Sleep(500 * time.Millisecond)
-
-		files, err = os.ReadDir(pidsDir)
-		if err != nil {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if len(files) != 2 {
-			t.Errorf("Expected 2 PID files, got %d", len(files))
-		}
+		waitForPids(t, pidsDir, 2)
 
 		// Wait for first process to exit naturally (triggering defer)
 		if err := cmd1.Wait(); err != nil {
 			t.Fatalf("cmd1 failed: %v", err)
 		}
-		time.Sleep(500 * time.Millisecond)
-
-		files, err = os.ReadDir(pidsDir)
-		if err != nil {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if len(files) != 1 {
-			t.Errorf("Expected 1 PID file after first exit, got %d", len(files))
+		
+		// Poll for the count to drop back to 1
+		timeout := time.After(2 * time.Second)
+		found := false
+		for !found {
+			select {
+			case <-timeout:
+				t.Fatalf("Timed out waiting for first process to unregister")
+			default:
+				files, err := os.ReadDir(pidsDir)
+				if err == nil && len(files) == 1 {
+					found = true
+					break
+				}
+				time.Sleep(50 * time.Millisecond)
+			}
 		}
 
 		// Wait for second process to exit naturally
 		if err := cmd2.Wait(); err != nil {
 			t.Fatalf("cmd2 failed: %v", err)
 		}
-		time.Sleep(500 * time.Millisecond)
-
-		// Verify a clean state
-		files, err = os.ReadDir(pidsDir)
-		if err != nil && !os.IsNotExist(err) {
-			t.Fatalf("Failed to read pids dir: %v", err)
-		}
-		if err == nil && len(files) != 0 {
-			t.Errorf("Expected 0 PID files after all exits, got %d", len(files))
+		
+		// Verify a clean state (expect 0 files)
+		timeout = time.After(2 * time.Second)
+		found = false
+		for !found {
+			select {
+			case <-timeout:
+				files, _ := os.ReadDir(pidsDir)
+				t.Fatalf("Expected 0 PID files after all exits, got %d", len(files))
+			default:
+				files, err := os.ReadDir(pidsDir)
+				if err != nil && os.IsNotExist(err) {
+					found = true
+					break
+				}
+				if err == nil && len(files) == 0 {
+					found = true
+					break
+				}
+				time.Sleep(50 * time.Millisecond)
+			}
 		}
 	})
 }