Allow a server api to implement a custom healthcheck (#109)

Author: mraerino

router/middleware.go

@@ -90,9 +90,7 @@ func HealthCheck(route string, f APIHandler) Middleware {
 				w.WriteHeader(http.StatusOK)
 				return
 			}
-			if err := f(w, r); err != nil {
-				HandleError(f(w, r), w, r)
-			}
+			HandleError(f(w, r), w, r)
 			return
 		}
 		next.ServeHTTP(w, r)

server/server.go

@@ -33,10 +33,22 @@ type APIDefinition interface {
 	Stop()
 }
 
+// HealthChecker is used to run a custom health check
+// Implement it on your API if you want it to be checked
+// when the healthcheck is called
+type HealthChecker interface {
+	Healthy(w http.ResponseWriter, r *http.Request) *router.HTTPError
+}
+
 func New(log logrus.FieldLogger, projectName string, config Config, api APIDefinition) (*Server, error) {
+	var healthHandler router.APIHandler
+	if checker, ok := api.(HealthChecker); ok {
+		healthHandler = checker.Healthy
+	}
+
 	r := router.New(
 		log,
-		router.OptHealthCheck(config.HealthPath, nil),
+		router.OptHealthCheck(config.HealthPath, healthHandler),
 		router.OptTracingMiddleware(log, projectName),
 	)
 

server/server_test.go

@@ -1,15 +1,16 @@
 package server
 
 import (
-	"github.com/netlify/netlify-commons/router"
-	"github.com/sirupsen/logrus"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"net/http"
 	"net/http/httptest"
 	"os"
 	"strings"
 	"testing"
+
+	"github.com/netlify/netlify-commons/router"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func init() {
@@ -42,6 +43,36 @@ func TestServerHealth(t *testing.T) {
 	assert.Equal(t, http.StatusOK, rsp.StatusCode)
 }
 
+type testAPICustomHealth struct{}
+
+func (a *testAPICustomHealth) Start(r router.Router) error {
+	r.Get("/", func(w http.ResponseWriter, r *http.Request) *router.HTTPError {
+		return nil
+	})
+	return nil
+}
+
+func (a *testAPICustomHealth) Stop() {}
+
+func (a *testAPICustomHealth) Healthy(w http.ResponseWriter, r *http.Request) *router.HTTPError {
+	return router.InternalServerError("healthcheck failed")
+}
+
+func TestServerCustomHealth(t *testing.T) {
+	apiDef := new(testAPICustomHealth)
+
+	cfg := testConfig()
+	svr, err := New(tl(t), "testing", cfg, apiDef)
+	require.NoError(t, err)
+
+	testSvr := httptest.NewServer(svr.svr.Handler)
+	defer testSvr.Close()
+
+	rsp, err := http.Get(testSvr.URL + cfg.HealthPath)
+	require.NoError(t, err)
+	assert.Equal(t, http.StatusInternalServerError, rsp.StatusCode)
+}
+
 func tl(t *testing.T) *logrus.Entry {
 	return logrus.WithField("test", t.Name())
 }