Refactor healthcheck logic into SetupHealthcheck in health.go and add integration test

endlessh_integration_test.go

@@ -263,7 +263,7 @@ func TestEndlesshIntegration_PrometheusMetrics(t *testing.T) {
 	cmd := exec.Command(
 		"go", "run", "main.go",
 		"-port=0",
-		"-enable_prometheus",
+		"-prometheus_enable",
 		"-prometheus_port=0",
 		"-interval_ms=100",
 		"-logtostderr", "-v=1",
@@ -322,3 +322,61 @@ func TestEndlesshIntegration_PrometheusMetrics(t *testing.T) {
 		t.Errorf("Expected bytes metric not found:\n%s", body)
 	}
 }
+
+func TestEndlesshIntegration_Healthcheck(t *testing.T) {
+	var stderr bytes.Buffer
+
+	cmd := exec.Command(
+		"go", "run", "main.go",
+		"-port=0",
+		"-healthcheck_enable",
+		"-healthcheck_port=0",
+		"-interval_ms=100",
+		"-logtostderr", "-v=1",
+	)
+	cmd.Stderr = &stderr
+
+	if err := cmd.Start(); err != nil {
+		t.Fatalf("Failed to start server: %v", err)
+	}
+	defer cmd.Process.Kill()
+
+	if !waitForLogMatch(&stderr, "Starting healthcheck on http", waitForListenTimeout) {
+		t.Fatalf("Healthcheck listener did not start: %s", stderr.String())
+	}
+
+	reHealth := regexp.MustCompile(`Starting healthcheck on http://.*:(\d+)/health`)
+	healthMatch := reHealth.FindStringSubmatch(stderr.String())
+	if len(healthMatch) < 2 {
+		t.Fatalf("Could not parse healthcheck port: %s", stderr.String())
+	}
+	healthPort := healthMatch[1]
+
+	// 1. Verify healthcheck HTTP endpoint returns valid JSON
+	resp, err := net.Dial("tcp", "localhost:"+healthPort)
+	if err != nil {
+		t.Fatalf("Failed to connect to healthcheck endpoint: %v", err)
+	}
+	fmt.Fprintf(resp, "GET /health HTTP/1.1\r\nHost: localhost\r\n\r\n")
+
+	buf := make([]byte, 8192)
+	n, _ := resp.Read(buf)
+	body := string(buf[:n])
+	resp.Close()
+
+	if !strings.Contains(body, `"status":"ok"`) {
+		t.Errorf("Expected status:ok in health response: %s", body)
+	}
+
+	// 2. Run the probe mode (-healthcheck) and verify it exits with 0
+	probeCmd := exec.Command(
+		"go", "run", "main.go",
+		"-healthcheck",
+		"-healthcheck_port="+healthPort,
+	)
+	var probeStderr bytes.Buffer
+	probeCmd.Stderr = &probeStderr
+	if err := probeCmd.Run(); err != nil {
+		t.Errorf("Healthcheck probe failed: %v, stderr: %s", err, probeStderr.String())
+	}
+}

health/health.go

@@ -21,6 +21,7 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"strconv"
 	"time"
 
 	"github.com/golang/glog"
@@ -84,3 +85,29 @@ func Probe(host, port string) bool {
 	}
 	return body.Status == "ok"
 }
+
+// SetupHealthcheck probes or starts the healthcheck listener as configured.
+func SetupHealthcheck(healthcheck bool, healthcheckEnabled bool, connType string, healthcheckHost, healthcheckPort *string) {
+	if healthcheck {
+		if !Probe(*healthcheckHost, *healthcheckPort) {
+			os.Exit(1)
+		}
+		os.Exit(0)
+	}
+
+	if healthcheckEnabled {
+		if connType == "tcp6" && *healthcheckHost == "0.0.0.0" {
+			*healthcheckHost = "[::]"
+		}
+		if *healthcheckPort == "0" || *healthcheckPort == "" {
+			l, err := net.Listen("tcp", *healthcheckHost+":0")
+			if err != nil {
+				glog.Fatalf("Failed to pick a free healthcheck port: %v", err)
+			}
+			actualPort := l.Addr().(*net.TCPAddr).Port
+			*healthcheckPort = strconv.Itoa(actualPort)
+			l.Close()
+		}
+		StartListener(*healthcheckHost, *healthcheckPort)
+	}
+}

main.go

@@ -167,12 +167,7 @@ func main() {
 	}
 	flag.Parse()
 
-	if *healthcheck {
+	health.SetupHealthcheck(*healthcheck, *healthcheckEnabled, *connType, healthcheckHost, healthcheckPort)
-		if !health.Probe(*healthcheckHost, *healthcheckPort) {
-			os.Exit(1)
-		}
-		os.Exit(0)
-	}
 
 	prometheusEnabled := *prometheusEnabledNew
 	prometheusEnableSet := false
@@ -208,21 +203,6 @@ func main() {
 		})
 	clients := startSending(*maxClients, *bannerMaxLength, records)
 
-	if *healthcheckEnabled {
-		if *connType == "tcp6" && *healthcheckHost == "0.0.0.0" {
-			*healthcheckHost = "[::]"
-		}
-		if *healthcheckPort == "0" || *healthcheckPort == "" {
-			l, err := net.Listen("tcp", *healthcheckHost+":0")
-			if err != nil {
-				glog.Fatalf("Failed to pick a free healthcheck port: %v", err)
-			}
-			actualPort := l.Addr().(*net.TCPAddr).Port
-			*healthcheckPort = strconv.Itoa(actualPort)
-			l.Close()
-		}
-		health.StartListener(*healthcheckHost, *healthcheckPort)
-	}
 
 	interval := time.Duration(*intervalMs) * time.Millisecond
 	// Listen for incoming connections.