-

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",

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()))

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

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())
 }