Support failing root handler on env var (#23)

Author: maruina

.pre-commit-config.yaml

@@ -1,10 +1,10 @@
+repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
   rev: v2.4.0
   hooks:
   - id: trailing-whitespace
   - id: end-of-file-fixer
-- repo: git://github.com/dnephin/pre-commit-golang
+- repo: https://github.com/golangci/golangci-lint
   rev: master
   hooks:
-    - id: go-fmt
-    - id: go-lint
+    - id: golangci-lint

README.md

@@ -18,6 +18,10 @@ To override the default values:
 
 * _INFRABIN_MAX_DELAY_ to change the maximum value for the `/delay` endpoint. Default to 120.
 
+## Environment variables
+
+* `FAIL_ROOT_HANDLER`: if set, the `/` endpoint will return a 503. This is useful when doing a B/G deployment to test the failure and rollback scenario.
+
 ## Service Endpoints
 
 * `GET /`

cmd/go-infrabin/main.go

@@ -14,26 +14,33 @@ import (
 
 // RootHandler handles the "/" endpoint
 func RootHandler(w http.ResponseWriter, r *http.Request) {
-	hostname, err := os.Hostname()
-	if err != nil {
-		log.Fatalf("cannot get hostname: %v", err)
-	}
 	w.Header().Set("Content-Type", "application/json")
-	w.WriteHeader(http.StatusOK)
 
-	var resp helpers.Response
-	resp.Hostname = hostname
-	resp.KubeResponse = &helpers.KubeResponse{
-		PodName:   helpers.GetEnv("POD_NAME", ""),
-		Namespace: helpers.GetEnv("POD_NAMESPACE", ""),
-		PodIP:     helpers.GetEnv("POD_IP", ""),
-		NodeName:  helpers.GetEnv("NODE_NAME", ""),
-	}
+	fail := helpers.GetEnv("FAIL_ROOT_HANDLER", "")
+	if fail != "" {
+		w.WriteHeader(http.StatusServiceUnavailable)
+	} else {
+		w.WriteHeader(http.StatusOK)
 
-	data := helpers.MarshalResponseToString(resp)
-	_, err = io.WriteString(w, data)
-	if err != nil {
-		log.Fatal("error writing to ResponseWriter: ", err)
+		hostname, err := os.Hostname()
+		if err != nil {
+			log.Fatalf("cannot get hostname: %v", err)
+		}
+
+		var resp helpers.Response
+		resp.Hostname = hostname
+		resp.KubeResponse = &helpers.KubeResponse{
+			PodName:   helpers.GetEnv("POD_NAME", ""),
+			Namespace: helpers.GetEnv("POD_NAMESPACE", ""),
+			PodIP:     helpers.GetEnv("POD_IP", ""),
+			NodeName:  helpers.GetEnv("NODE_NAME", ""),
+		}
+
+		data := helpers.MarshalResponseToString(resp)
+		_, err = io.WriteString(w, data)
+		if err != nil {
+			log.Fatal("error writing to ResponseWriter: ", err)
+		}
 	}
 }
 

cmd/go-infrabin/main_test.go

@@ -44,6 +44,28 @@ func TestRootHandler(t *testing.T) {
 	}
 }
 
+func TestFailRootHandler(t *testing.T) {
+	if err := os.Setenv("FAIL_ROOT_HANDLER", "true"); err != nil {
+		t.Errorf("cannot set environment variable")
+	}
+	req, err := http.NewRequest("GET", "/", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rr := httptest.NewRecorder()
+	handler := http.HandlerFunc(RootHandler)
+	handler.ServeHTTP(rr, req)
+
+	if status := rr.Code; status != http.StatusServiceUnavailable {
+		t.Errorf("handler returned wrong status code: got %v want %v",
+			status, http.StatusServiceUnavailable)
+	}
+	if err = os.Unsetenv("FAIL_ROOT_HANDLER"); err != nil {
+		t.Fatal(err)
+	}
+}
+
 func TestRootHandlerKubernetes(t *testing.T) {
 	// Set Kubernetes OS env variables
 	podName := "go-infrabin-hjv8k"

internal/helpers/struct_test.go

@@ -1,7 +1,6 @@
 package helpers
 
 import (
-	"fmt"
 	"log"
 	"testing"
 )
@@ -11,7 +10,7 @@ func TestMarshalResponseToString(t *testing.T) {
 	resp.Hostname = "hal"
 
 	data := MarshalResponseToString(resp)
-	expected := fmt.Sprint(`{"hostname":"hal"}`)
+	expected := `{"hostname":"hal"}`
 	if data != expected {
 		log.Fatalf("error: expected %v got %v", expected, data)
 	}