1 files changed, 25 insertions, 2 deletions
diff --git a/internal/namespace/lifecycle_test.go b/internal/namespace/lifecycle_test.go
index 1fb0a13..9962e14 100644
--- a/internal/namespace/lifecycle_test.go
+++ b/internal/namespace/lifecycle_test.go
@@ -87,17 +87,40 @@ func TestLifecycleReferenceCounting(t *testing.T) {
 		if err := RegisterProcess(pm, profile); err != nil {
 			t.Fatal(err)
 		}
+		
+		// Simulate the application flow: Unregister before checking if we are the last one
+		if err := UnregisterProcess(pm, profile); err != nil {
+			t.Fatal(err)
+		}
+
 		isLast, err = IsLastProcess(pm, profile)
 		if err != nil || !isLast {
-			t.Errorf("Expected IsLastProcess to be true for single process, got %v, err: %v", isLast, err)
+			t.Errorf("Expected IsLastProcess to be true after unregistering the only process, got %v, err: %v", isLast, err)
 		}
 
+		// Add a "stale" process to ensure it's pruned and doesn't count as active
 		if err := os.WriteFile(filepath.Join(pidsDir, "1234567"), []byte(""), 0644); err != nil {
 			t.Fatal(err)
 		}
+		
+		// Register a real process so that pruning has something to do
+		if err := RegisterProcess(pm, profile); err != nil {
+			t.Fatal(err)
+		}
+		
+		// Prune the stale one
+		if err := PruneStalePids(pm, profile); err != nil {
+			t.Fatal(err)
+		}
+
+		// Unregister the real one
+		if err := UnregisterProcess(pm, profile); err != nil {
+			t.Fatal(err)
+		}
+
 		isLast, err = IsLastProcess(pm, profile)
 		if err != nil || !isLast {
-			t.Errorf("Expected IsLastProcess to be true because 1234567 is dead, got %v, err: %v", isLast, err)
+			t.Errorf("Expected IsLastProcess to be true after pruning stale and unregistering current, got %v, err: %v", isLast, err)
 		}
 	})
 }