From 875381a840456e3559ae52a8b38d61e0c3af1156 Mon Sep 17 00:00:00 2001
From: Vincent Bernat <vincent@bernat.im>
Date: Sun, 29 Jul 2018 18:08:11 +0200
Subject: [PATCH] healthcheck: implement GetOrRegisterHealthcheck and
 NewRegisteredHealthcheck

Other metric implementations have those methods. Let's add them for
healthchecks too.
---
 healthcheck.go      | 19 +++++++++++++++++++
 healthcheck_test.go | 17 +++++++++++++++++
 2 files changed, 36 insertions(+)
 create mode 100644 healthcheck_test.go

diff --git a/healthcheck.go b/healthcheck.go
index 445131c..0f61b95 100644
--- a/healthcheck.go
+++ b/healthcheck.go
@@ -8,6 +8,15 @@ type Healthcheck interface {
 	Unhealthy(error)
 }
 
+// GetOrRegisterHealthcheck returns an existing Healthcheck or
+// constructs and registers a new StandardHealthcheck.
+func GetOrRegisterHealthcheck(name string, r Registry, f func(Healthcheck)) Healthcheck {
+	if nil == r {
+		r = DefaultRegistry
+	}
+	return r.GetOrRegister(name, func() Healthcheck { return NewHealthcheck(f) }).(Healthcheck)
+}
+
 // NewHealthcheck constructs a new Healthcheck which will use the given
 // function to update its status.
 func NewHealthcheck(f func(Healthcheck)) Healthcheck {
@@ -17,6 +26,16 @@ func NewHealthcheck(f func(Healthcheck)) Healthcheck {
 	return &StandardHealthcheck{nil, f}
 }
 
+// NewRegisteredHealthcheck constructs and registers a new StandardHealthcheck.
+func NewRegisteredHealthcheck(name string, r Registry, f func(Healthcheck)) Healthcheck {
+	c := NewHealthcheck(f)
+	if nil == r {
+		r = DefaultRegistry
+	}
+	r.Register(name, c)
+	return c
+}
+
 // NilHealthcheck is a no-op.
 type NilHealthcheck struct{}
 
diff --git a/healthcheck_test.go b/healthcheck_test.go
new file mode 100644
index 0000000..9ad4246
--- /dev/null
+++ b/healthcheck_test.go
@@ -0,0 +1,17 @@
+package metrics
+
+import (
+	"errors"
+	"testing"
+)
+
+func TestGetOrRegisterHealthcheck(t *testing.T) {
+	r := NewRegistry()
+	check := func(h Healthcheck) {
+		h.Unhealthy(errors.New("foo"))
+	}
+	NewRegisteredHealthcheck("foo", r, check).Check()
+	if h := GetOrRegisterHealthcheck("foo", r, check); h.Error().Error() != "foo" {
+		t.Fatal(h)
+	}
+}