urlchecker with statuscode, bodystring and http method

Author: mschneider82

url/checker.go

@@ -1,6 +1,9 @@
 package url
 
 import (
+	"bytes"
+	"context"
+	"io/ioutil"
 	"net/http"
 	"time"
 
@@ -9,13 +12,16 @@ import (
 
 // Checker is a checker that check a given URL
 type Checker struct {
-	URL     string
-	Timeout time.Duration
+	URL                string
+	Timeout            time.Duration
+	Method             string // Defaults to GET
+	ExpectedBodyString string // Disabled if empty
+	ExpectedStatusCode int    // if 0 defaults to 200
 }
 
 // NewChecker returns a new url.Checker with the given URL
 func NewChecker(url string) Checker {
-	return Checker{URL: url, Timeout: 5 * time.Second}
+	return NewCheckerWithTimeout(url, 5*time.Second)
 }
 
 // NewCheckerWithTimeout returns a new url.Checker with the given URL and given timeout
@@ -27,31 +33,46 @@ func NewCheckerWithTimeout(url string, timeout time.Duration) Checker {
 // If the request returns something different than StatusOK,
 // The status is set to StatusBadRequest and the URL is considered Down
 func (u Checker) Check() health.Health {
-	client := http.Client{
-		Timeout: u.Timeout,
-	}
-
+	ctxWithTimeout, cancelFn := context.WithTimeout(context.Background(), u.Timeout)
+	defer cancelFn()
 	health := health.NewHealth()
+	body := new(bytes.Buffer)
+	req, err := http.NewRequestWithContext(ctxWithTimeout, u.Method, u.URL, body)
+	if err != nil {
+		health.Down().AddInfo("code", http.StatusBadRequest).AddInfo("error", err.Error())
+		return health
+	}
 
-	resp, err := client.Head(u.URL)
-
+	resp, err := http.DefaultClient.Do(req)
 	if resp != nil {
 		defer resp.Body.Close()
 	}
-
 	if err != nil {
-		health.Down().AddInfo("code", http.StatusBadRequest)
-
+		health.Down().AddInfo("code", http.StatusBadRequest).AddInfo("error", err.Error())
 		return health
 	}
 
-	if resp.StatusCode == http.StatusOK {
-		health.Up()
+	if u.ExpectedStatusCode == 0 {
+		u.ExpectedStatusCode = http.StatusOK
+	}
+
+	if resp.StatusCode == u.ExpectedStatusCode {
+		health.Up().AddInfo("code", resp.StatusCode)
 	} else {
-		health.Down()
+		health.Down().AddInfo("code", resp.StatusCode).AddInfo("expectedCode", u.ExpectedStatusCode)
+		return health
 	}
 
-	health.AddInfo("code", resp.StatusCode)
+	if u.ExpectedBodyString != "" {
+		respbody, err := ioutil.ReadAll(resp.Body)
+		if err != nil {
+			health.Down().AddInfo("error", err.Error())
+			return health
+		}
+		if !bytes.ContainsAny(respbody, u.ExpectedBodyString) {
+			health.Down().AddInfo("body", "does not contain: "+u.ExpectedBodyString)
+		}
+	}
 
 	return health
 }

url/checker_test.go

@@ -49,6 +49,52 @@ func Test_Checker_Check_Up(t *testing.T) {
 	check(t, resp, wants, http.StatusOK)
 }
 
+func Test_Checker_Check_Up_BodyCheck(t *testing.T) {
+	mux := http.NewServeMux()
+
+	server := httptest.NewServer(mux)
+
+	checker := NewChecker(fmt.Sprintf("%s/up/", server.URL))
+	checker.ExpectedBodyString = "pong"
+
+	handler := health.NewHandler()
+	handler.AddChecker("Up", checker)
+
+	mux.Handle("/health/", handler)
+	mux.HandleFunc("/up/", func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, "pong")
+	})
+
+	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))
+
+	wants := `{"Up":{"code":200,"status":"UP"},"status":"UP"}`
+
+	check(t, resp, wants, http.StatusOK)
+}
+
+func Test_Checker_Check_Down_BodyCheck(t *testing.T) {
+	mux := http.NewServeMux()
+
+	server := httptest.NewServer(mux)
+
+	checker := NewChecker(fmt.Sprintf("%s/up/", server.URL))
+	checker.ExpectedBodyString = "pong"
+
+	handler := health.NewHandler()
+	handler.AddChecker("Down", checker)
+
+	mux.Handle("/health/", handler)
+	mux.HandleFunc("/up/", func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, "fail")
+	})
+
+	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))
+
+	wants := `{"Down":{"body":"does not contain: pong","code":200,"status":"DOWN"},"status":"DOWN"}`
+
+	check(t, resp, wants, http.StatusServiceUnavailable)
+}
+
 func Test_Checker_Check_Down(t *testing.T) {
 	mux := http.NewServeMux()
 
@@ -67,7 +113,7 @@ func Test_Checker_Check_Down(t *testing.T) {
 
 	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))
 
-	wants := `{"Down":{"code":500,"status":"DOWN"},"status":"DOWN"}`
+	wants := `{"Down":{"code":500,"expectedCode":200,"status":"DOWN"},"status":"DOWN"}`
 
 	check(t, resp, wants, http.StatusServiceUnavailable)
 }
@@ -86,7 +132,7 @@ func Test_Checker_Check_Down_invalid(t *testing.T) {
 
 	resp, _ := http.Get(fmt.Sprintf("%s/health/", server.URL))
 
-	wants := `{"Down":{"code":400,"status":"DOWN"},"status":"DOWN"}`
+	wants := `{"Down":{"code":400,"error":"Get : unsupported protocol scheme \"\"","status":"DOWN"},"status":"DOWN"}`
 	check(t, resp, wants, http.StatusServiceUnavailable)
 }