Merge pull request #99 from NETWAYS/feature/health-crit

Add CLI flag to override exit code if API is unreachable

Author: martialblog

README.md

@@ -61,6 +61,7 @@ Flags:
       --heap-usage-threshold-crit string        The percentage relative to the heap size limit on which to be a critical result (default "80")
       --cpu-usage-threshold-warn string         The percentage of CPU usage on which to be a warning result (default "100")
       --cpu-usage-threshold-crit string         The percentage of CPU usage on which to be a critical result (default "100")
+      --unreachable-state int                   Exit with specified code if unreachable. Examples: 1 for Warning, 2 for Critical, 3 for Unknown (default 3)
   -h, --help                                    help for health
 ```
 

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. Examples: 1 for Warning, 2 for Critical, 3 for Unknown")
+
 	fs.SortFlags = false
 }

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