fix test

sreya · sreya · commit 1703e6d6c406 · 2022-06-16T21:29:28.000Z
diff --git a/agent/reaper/reaper.go b/agent/reaper/reaper.go
@@ -13,7 +13,7 @@ import (
 // to indicated that the process is a child as opposed to the reaper.
 // Since we are forkexec'ing we need to be able to differentiate between
 // the two to avoid fork bombing ourselves.
-const agentEnvMark = "CODER_AGENT"
+const agentEnvMark = "CODER_REAPER_AGENT"
 
 // IsChild returns true if we're the forked process.
 func IsChild() bool {
@@ -29,6 +29,8 @@ func IsInitProcess() bool {
 // complications with spawning `exec.Commands` in the same process that
 // is reaping, we forkexec a child process. This prevents a race between
 // the reaper and an exec.Command waiting for its process to complete.
+// The provided 'pids' channel may be nil if the caller does not care about the
+// reaped children PIDs.
 func ForkReap(pids reap.PidCh) error {
 	// Check if the process is the parent or the child.
 	// If it's the child we want to skip attempting to reap.
diff --git a/agent/reaper/reaper_test.go b/agent/reaper/reaper_test.go
@@ -1,6 +1,7 @@
 package reaper_test
 
 import (
+	"fmt"
 	"os/exec"
 	"testing"
 	"time"
@@ -19,6 +20,9 @@ func TestReap(t *testing.T) {
 		t.Skip("I'm a child!")
 	}
 
+	// OK checks that's the reaper is successfully reaping
+	// exited processes and passing the PIDs through the shared
+	// channel.
 	t.Run("OK", func(t *testing.T) {
 		pids := make(reap.PidCh, 1)
 		err := reaper.ForkReap(pids)
@@ -41,11 +45,14 @@ func TestReap(t *testing.T) {
 		expectedPIDs := []int{cmd.Process.Pid, cmd2.Process.Pid}
 
 		deadline := time.NewTimer(time.Second * 5)
-		select {
-		case <-deadline.C:
-			t.Fatalf("Timed out waiting for process")
-		case pid := <-pids:
-			require.Contains(t, expectedPIDs, pid)
+		for i := 0; i < len(expectedPIDs); i++ {
+			select {
+			case <-deadline.C:
+				t.Fatalf("Timed out waiting for process")
+			case pid := <-pids:
+				fmt.Println("pid: ", pid)
+				require.Contains(t, expectedPIDs, pid)
+			}
 		}
 	})
 }
