Add health command for hub and align agent health command

beszel/cmd/agent/agent.go

@@ -6,7 +6,6 @@ import (
 	"flag"
 	"fmt"
 	"log"
-	"log/slog"
 	"os"
 
 	"golang.org/x/crypto/ssh"
@@ -56,9 +55,12 @@ func (opts *cmdOptions) parse() bool {
 		flag.CommandLine.Parse(args)
 		addr := opts.getAddress()
 		network := agent.GetNetwork(addr)
-		exitCode, err := agent.Health(addr, network)
-		slog.Info("Health", "code", exitCode, "err", err)
-		os.Exit(exitCode)
+		err := agent.Health(addr, network)
+		if err != nil {
+			log.Fatal(err)
+		}
+		fmt.Print("ok")
+		return true
 	}
 
 	flag.Parse()

beszel/cmd/hub/hub.go

@@ -4,7 +4,10 @@ import (
 	"beszel"
 	"beszel/internal/hub"
 	_ "beszel/migrations"
+	"fmt"
+	"net/http"
 	"os"
+	"time"
 
 	"github.com/pocketbase/pocketbase"
 	"github.com/pocketbase/pocketbase/plugins/migratecmd"
@@ -12,6 +15,16 @@ import (
 )
 
 func main() {
+	// handle health check first to prevent unneeded execution
+	if len(os.Args) > 3 && os.Args[1] == "health" {
+		url := os.Args[3]
+		if err := checkHealth(url); err != nil {
+			log.Fatal(err)
+		}
+		fmt.Print("ok")
+		return
+	}
+
 	baseApp := getBaseApp()
 	h := hub.NewHub(baseApp)
 	h.BootstrapHub()
@@ -35,6 +48,8 @@ func getBaseApp() *pocketbase.PocketBase {
 		Short: "Update " + beszel.AppName + " to the latest version",
 		Run:   hub.Update,
 	})
+	// add health command
+	baseApp.RootCmd.AddCommand(newHealthCmd())
 
 	// enable auto creation of migration files when making collection changes in the Admin UI
 	migratecmd.MustRegister(baseApp, baseApp.RootCmd, migratecmd.Config{
@@ -44,3 +59,39 @@ func getBaseApp() *pocketbase.PocketBase {
 
 	return baseApp
 }
+
+func newHealthCmd() *cobra.Command {
+	var baseURL string
+
+	healthCmd := &cobra.Command{
+		Use:   "health",
+		Short: "Check health of running hub",
+		Run: func(cmd *cobra.Command, args []string) {
+			if err := checkHealth(baseURL); err != nil {
+				log.Fatal(err)
+			}
+			os.Exit(0)
+		},
+	}
+	healthCmd.Flags().StringVar(&baseURL, "url", "", "base URL")
+	healthCmd.MarkFlagRequired("url")
+	return healthCmd
+}
+
+// checkHealth checks the health of the hub.
+func checkHealth(baseURL string) error {
+	client := &http.Client{
+		Timeout: time.Second * 3,
+	}
+	healthURL := baseURL + "/api/health"
+	resp, err := client.Get(healthURL)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		return fmt.Errorf("%s returned status %d", healthURL, resp.StatusCode)
+	}
+	return nil
+}

beszel/internal/agent/health.go

@@ -6,14 +6,13 @@ import (
 )
 
 // Health checks if the agent's server is running by attempting to connect to it.
-// It returns 0 if the server is running, 1 otherwise (as in exit codes).
 //
 // If an error occurs when attempting to connect to the server, it returns the error.
-func Health(addr string, network string) (int, error) {
+func Health(addr string, network string) error {
 	conn, err := net.DialTimeout(network, addr, 4*time.Second)
 	if err != nil {
-		return 1, err
+		return err
 	}
 	conn.Close()
-	return 0, nil
+	return nil
 }

beszel/internal/agent/health_test.go

@@ -9,7 +9,6 @@ import (
 	"os"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"beszel/internal/agent"
@@ -77,11 +76,8 @@ func TestHealth(t *testing.T) {
 		defer cleanup()
 
 		// Run the health check with explicit parameters
-		result, err := agent.Health(sockFile, "unix")
+		err := agent.Health(sockFile, "unix")
 		require.NoError(t, err, "Failed to check health")
-
-		// Verify the result
-		assert.Equal(t, 0, result, "Health check should return 0 when server is running")
 	})
 
 	t.Run("server is running (tcp address)", func(t *testing.T) {
@@ -90,11 +86,8 @@ func TestHealth(t *testing.T) {
 		defer cleanup()
 
 		// Run the health check with explicit parameters
-		result, err := agent.Health(addr, "tcp")
+		err := agent.Health(addr, "tcp")
 		require.NoError(t, err, "Failed to check health")
-
-		// Verify the result
-		assert.Equal(t, 0, result, "Health check should return 0 when server is running")
 	})
 
 	t.Run("server is not running", func(t *testing.T) {
@@ -102,18 +95,14 @@ func TestHealth(t *testing.T) {
 		addr := "127.0.0.1:65535"
 
 		// Run the health check with explicit parameters
-		result, err := agent.Health(addr, "tcp")
+		err := agent.Health(addr, "tcp")
 		require.Error(t, err, "Health check should return an error when server is not running")
-
-		// Verify the result
-		assert.Equal(t, 1, result, "Health check should return 1 when server is not running")
 	})
 
 	t.Run("invalid network", func(t *testing.T) {
 		// Use an invalid network type
-		result, err := agent.Health("127.0.0.1:8080", "invalid_network")
+		err := agent.Health("127.0.0.1:8080", "invalid_network")
 		require.Error(t, err, "Health check should return an error with invalid network")
-		assert.Equal(t, 1, result, "Health check should return 1 when network is invalid")
 	})
 
 	t.Run("unix socket not found", func(t *testing.T) {
@@ -123,8 +112,7 @@ func TestHealth(t *testing.T) {
 		// Make sure it really doesn't exist
 		os.Remove(nonExistentSocket)
 
-		result, err := agent.Health(nonExistentSocket, "unix")
+		err := agent.Health(nonExistentSocket, "unix")
 		require.Error(t, err, "Health check should return an error when socket doesn't exist")
-		assert.Equal(t, 1, result, "Health check should return 1 when socket doesn't exist")
 	})
 }