From 1d9f971fa3246ef19a5ec290751c0bda9c4c1786 Mon Sep 17 00:00:00 2001
From: Markus Opolka <markus.opolka@netways.de>
Date: Mon, 19 Feb 2024 20:35:13 +0100
Subject: [PATCH] Add CLI flag to override exit code if API is unreachable

---
 cmd/health.go      |  6 +++++-
 cmd/health_test.go | 23 ++++++++++++++++++++++-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/cmd/health.go b/cmd/health.go
index 3c80520..4b77960 100644
--- a/cmd/health.go
+++ b/cmd/health.go
@@ -22,6 +22,7 @@ type HealthConfig struct {
 	HeapUseThresCritical  string
 	CPUUseThresWarning    string
 	CPUUseThresCritical   string
+	UnreachableExitCode   int
 }
 
 // To store the parsed CLI parameters.
@@ -164,7 +165,7 @@ var healthCmd = &cobra.Command{
 		resp, err := c.Client.Get(u)
 
 		if err != nil {
-			check.ExitError(err)
+			check.ExitRaw(cliHealthConfig.UnreachableExitCode, err.Error())
 		}
 
 		if resp.StatusCode != http.StatusOK {
@@ -284,5 +285,8 @@ func init() {
 	fs.StringVarP(&cliHealthConfig.CPUUseThresCritical, "cpu-usage-threshold-crit", "", "100",
 		"The percentage of CPU usage on which to be a critical result")
 
+	fs.IntVarP(&cliHealthConfig.UnreachableExitCode, "unreachable-state", "", 3,
+		"Exit with specified code if unreachable")
+
 	fs.SortFlags = false
 }
diff --git a/cmd/health_test.go b/cmd/health_test.go
index d559220..738ed73 100644
--- a/cmd/health_test.go
+++ b/cmd/health_test.go
@@ -10,7 +10,6 @@ import (
 )
 
 func TestHealth_ConnectionRefused(t *testing.T) {
-
 	cmd := exec.Command("go", "run", "../main.go", "health", "--port", "9999")
 	out, _ := cmd.CombinedOutput()
 
@@ -22,6 +21,28 @@ func TestHealth_ConnectionRefused(t *testing.T) {
 	}
 }
 
+func TestHealth_ConnectionRefusedCritical(t *testing.T) {
+	cmd := exec.Command("go", "run", "../main.go", "health", "--port", "9999", "--unreachable-state", "2")
+	out, _ := cmd.CombinedOutput()
+
+	actual := string(out)
+	expected := "[CRITICAL] - Get \"http://localhost:9999/"
+
+	if !strings.Contains(actual, expected) {
+		t.Error("\nActual: ", actual, "\nExpected: ", expected)
+	}
+
+	cmd = exec.Command("go", "run", "../main.go", "health", "--port", "9999", "--unreachable-state", "-123")
+	out, _ = cmd.CombinedOutput()
+
+	actual = string(out)
+	expected = "[UNKNOWN] - Get \"http://localhost:9999/"
+
+	if !strings.Contains(actual, expected) {
+		t.Error("\nActual: ", actual, "\nExpected: ", expected)
+	}
+}
+
 type HealthTest struct {
 	name     string
 	server   *httptest.Server