From 403c1b2a0c34de06c246b3442fb21c18e54ee5ce Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Mon, 17 Feb 2025 23:02:17 +0100
Subject: [PATCH] -

---
 cmd/selftest/child.go           |  1 +
 libs/daemon/daemon.go           | 13 ++++++--
 libs/process/testdata/parent.sh |  6 ++++
 libs/process/wait_test.go       | 58 +++++++++++++++++++++++++++++++--
 4 files changed, 73 insertions(+), 5 deletions(-)
 create mode 100755 libs/process/testdata/parent.sh

diff --git a/cmd/selftest/child.go b/cmd/selftest/child.go
index ebb319d34..b1584970b 100644
--- a/cmd/selftest/child.go
+++ b/cmd/selftest/child.go
@@ -12,6 +12,7 @@ import (
 	"github.com/spf13/cobra"
 )
 
+// TODO: Manually test that indeed latency is not added.
 func newChildCommand() *cobra.Command {
 	return &cobra.Command{
 		Use: "child",
diff --git a/libs/daemon/daemon.go b/libs/daemon/daemon.go
index 1c7dc7e9a..89b66e952 100644
--- a/libs/daemon/daemon.go
+++ b/libs/daemon/daemon.go
@@ -17,8 +17,10 @@ type Daemon struct {
 	// Environment variables to set in the child process.
 	Env []string
 
-	// Arguments to pass to the child process. The main executable is always the CLI
-	// binary itself.
+	// Path to executable to run. If empty, the current executable is used.
+	Executable string
+
+	// Arguments to pass to the child process.
 	Args []string
 
 	// Log file to write the child process's output to.
@@ -34,7 +36,12 @@ func (d *Daemon) Start() error {
 		return err
 	}
 
-	d.cmd = exec.Command(cli, d.Args...)
+	executable := d.Executable
+	if executable == "" {
+		executable = cli
+	}
+
+	d.cmd = exec.Command(executable, d.Args...)
 
 	// Set environment variable so that the child process know's it's parent's PID.
 	d.Env = append(d.Env, fmt.Sprintf("%s=%d", DatabricksCliParentPid, os.Getpid()))
diff --git a/libs/process/testdata/parent.sh b/libs/process/testdata/parent.sh
new file mode 100755
index 000000000..be80c0ee4
--- /dev/null
+++ b/libs/process/testdata/parent.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+# Launch background process.
+(sleep 5; echo "abc" > $2) &
+
+# Save PID of the background process to the file specified by the first argument.
+echo -n $! > $1
diff --git a/libs/process/wait_test.go b/libs/process/wait_test.go
index bc9407a8e..8db9de781 100644
--- a/libs/process/wait_test.go
+++ b/libs/process/wait_test.go
@@ -1,12 +1,66 @@
 package process
 
 import (
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strconv"
 	"testing"
+	"time"
 
+	"github.com/databricks/cli/internal/testutil"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestWait(t *testing.T) {
-	err := Wait(1000000)
-	assert.EqualError(t, err, "process with pid 1000000 does not exist")
+	t.Parallel()
+	tmpDir := t.TempDir()
+
+	pidFile := filepath.Join(tmpDir, "child.pid")
+	outputFile := filepath.Join(tmpDir, "output.txt")
+
+	// For this test, we cannot start the background process in this test itself
+	// and instead have to use parent.sh as an intermediary.
+	//
+	// This is because in Unix if we start the background process in this test itself,
+	// the background process will be a child of this test process and thus would
+	// need to be reaped by this test process (using the Wait function / syscall).
+	// Otherwise waitForPid will forever wait for the background process to finish.
+	//
+	// If we rely on an intermediate script to start the background process, the
+	// background process is reasigned to the init process (PID 1) once the parent
+	// exits and thus we can successfully wait for it in this test using waitForPid function.
+	cmd := exec.Command("./testdata/parent.sh", pidFile, outputFile)
+	err := cmd.Start()
+	require.NoError(t, err)
+
+	// Wait 5 seconds for the parent bash script to write the child's PID to the file.
+	var childPid int
+	require.Eventually(t, func() bool {
+		b, err := os.ReadFile(pidFile)
+		if err != nil {
+			return false
+		}
+
+		childPid, err = strconv.Atoi(string(b))
+		require.NoError(t, err)
+		return true
+	}, 2*time.Second, 100*time.Millisecond)
+
+	// The output file should not exist yet since the background process should
+	// still be running.
+	assert.NoFileExists(t, outputFile)
+
+	// Wait for the background process to finish.
+	err = waitForPid(childPid)
+	assert.NoError(t, err)
+
+	// The output file should exist now since the background process has finished.
+	testutil.AssertFileContents(t, outputFile, "abc\n")
+
+	// Since the background process has finished, waiting for it again should
+	// return an error.
+	err = waitForPid(childPid)
+	assert.Regexp(t, "process with pid .* does not exist", err.Error())
 }