Simplify and fix test.

Author: David Grieser

cmd/root.go

@@ -3,7 +3,7 @@ package cmd
 import (
 	"net"
 	"net/http"
-	"net/http/pprof"
+	_ "net/http/pprof"
 
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -25,21 +25,14 @@ var rootCmd = &cobra.Command{
 	PersistentPreRun: func(cmd *cobra.Command, args []string) {
 		if enableProfile {
 			go func() {
-				mux := http.NewServeMux()
-				mux.HandleFunc("/debug/pprof/", pprof.Index)
-				mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
-				mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
-				mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
-				mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
-
-				listener, err := net.Listen("tcp", ":0")
+				// pprof handlers are auto-registered on the default ServeMux when imported.
+				listener, err := net.Listen("tcp", "127.0.0.1:")
 				if err != nil {
 					log.Errorf("pprof server failed to listen: %v", err)
 					return
 				}
-				log.Infof("Starting pprof server on http://localhost%s/debug/pprof/", listener.Addr().String())
-				err = http.Serve(listener, mux)
-				if err != nil {
+				log.Infof("Starting pprof server on http://%s/debug/pprof/", listener.Addr().String())
+				if err = http.Serve(listener, nil); err != nil {
 					log.Errorf("pprof server error: %v", err)
 				}
 			}()

cmd/root_test.go

@@ -41,7 +41,7 @@ func TestProfileFlag(t *testing.T) {
 	hclContent := `
 job "dummy" {
   command = "/bin/sh"
-  args    = ["-c", "echo dummy job running; sleep 30"] // Increased sleep to 30s
+  args    = ["-c", "echo dummy job running; sleep 10"]
 }
 `
 	if err := os.WriteFile(filepath.Join(jobsDir, "dummy.hcl"), []byte(hclContent), 0644); err != nil {
@@ -105,7 +105,7 @@ job "dummy" {
 	go func() {
 		defer wg.Done()
 		reader := bufio.NewReader(stderrPipe) // Read from stderrPipe
-		pprofLogRegex := regexp.MustCompile(`Starting pprof server on http://localhost(?:[^:]+)?:(\d+)/debug/pprof/`)
+		pprofLogRegex := regexp.MustCompile(`Starting pprof server on http://127.0.0.1:(\d+)/debug/pprof/`)
 		for {
 			line, err := reader.ReadString('\n')
 			if len(line) > 0 {
@@ -141,11 +141,11 @@ job "dummy" {
 		t.Logf("[%s] Successfully found pprof server URL: %s", time.Since(startTime), pprofURL)
 	case <-ctx.Done(): // This is the main test context
 		t.Logf("[%s] Main context done while waiting for pprof URL.", time.Since(startTime))
-		cmd.Process.Kill() // Ensure process is killed if context times out
-		wg.Wait()          // Wait for goroutine to finish
+		cmd.Process.Kill()            // Ensure process is killed if context times out
+		wg.Wait()                     // Wait for goroutine to finish
 		logContent := errBuf.String() // Use the full stderr buffer for logging
 		t.Fatalf("Test timed out waiting for pprof server log. Stderr:\n%s", logContent)
-	case <-time.After(15 * time.Second): // Specific timeout for finding the log line, increased slightly
+	case <-time.After(5 * time.Second): // Specific timeout for finding the log line
 		t.Logf("[%s] Timed out waiting for pprof log line via channel.", time.Since(startTime))
 		cmd.Process.Kill() // Ensure process is killed
 		wg.Wait()          // Wait for goroutine to finish
@@ -217,7 +217,7 @@ job "dummy" {
 
 	if !success {
 		finalErrorDescriptive := "No successful HTTP GET." // Renamed to avoid conflict if needed, or ensure proper scope
-		if httpErr != nil { // httpErr should hold the error from the last attempt
+		if httpErr != nil {                                // httpErr should hold the error from the last attempt
 			finalErrorDescriptive = httpErr.Error()
 		}
 		statusCode := 0
@@ -244,7 +244,7 @@ job "dummy" {
 	// HTTP Check related variables (resp, httpErr, success) are already defined before this block.
 	// The defer func above will handle cmd.Wait() and final process state logging.
 	// We only need to kill the process here if the HTTP checks complete *successfully*
-	// and the process hasn't naturally exited due to its own short lifecycle (unlikely with sleep 30).
+	// and the process hasn't naturally exited due to its own short lifecycle (unlikely with sleep 10).
 	if success { // If HTTP check was successful
 		if cmd.Process != nil && (cmd.ProcessState == nil || !cmd.ProcessState.Exited()) {
 			t.Logf("[%s] HTTP check successful, ensuring mittnite process (PID: %d) is terminated.", time.Since(startTime), cmd.Process.Pid)