1 files changed, 25 insertions, 36 deletions
diff --git a/tests/e2e/lifecycle_test.go b/tests/e2e/lifecycle_test.go
index 0f6cae1..4b452da 100644
--- a/tests/e2e/lifecycle_test.go
+++ b/tests/e2e/lifecycle_test.go
@@ -37,6 +37,28 @@ func waitForPids(t *testing.T, pidsDir string, expectedCount int) {
 	}
 }
 
+func waitForLifecycle(t *testing.T, binaryPath, runtimeDir, profile string, expectedActive bool) {
+	timeout := time.After(2 * time.Second)
+	tick := time.NewTicker(50 * time.Millisecond)
+	defer tick.Stop()
+
+	for {
+		select {
+		case <-timeout:
+			t.Fatalf("Timed out waiting for lifecycle state: expected active=%v", expectedActive)
+		case <-tick.C:
+			cmd := exec.Command(binaryPath, "--profile", profile, "test-lifecycle")
+			cmd.Env = append(os.Environ(), fmt.Sprintf("XDG_RUNTIME_DIR=%s", runtimeDir))
+			err := cmd.Run()
+
+			isActive := err == nil
+			if isActive == expectedActive {
+				return
+			}
+		}
+	}
+}
+
 func TestNamespaceLifecycleAutomation(t *testing.T) {
 	// 1. Setup Environment
 	binaryPath, err := GetBinaryPath()
@@ -63,7 +85,7 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 		}
 
 		// Verify PID file exists using polling
-		waitForPids(t, pidsDir, 1)
+		waitForLifecycle(t, binaryPath, tmpRuntimeDir, "default", true)
 
 		// Start a second process using the same profile
 		cmd2 := exec.Command(binaryPath, "--profile", "default", "--", "sleep", "0.1")
@@ -79,21 +101,7 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 		}
 
 		// 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)
-			}
-		}
+		waitForLifecycle(t, binaryPath, tmpRuntimeDir, "default", true)
 
 		// Wait for second process to exit naturally
 		if err := cmd2.Wait(); err != nil {
@@ -101,25 +109,6 @@ func TestNamespaceLifecycleAutomation(t *testing.T) {
 		}
 
 		// 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)
-			}
-		}
+		waitForLifecycle(t, binaryPath, tmpRuntimeDir, "default", false)
 	})
 }