Check framework (#32)

* Check rewrite using new basic Health Check framework

* Removing unused method

Co-authored-by: Shaun Davis <[email protected]>

Author: davissp14

go.mod

@@ -3,6 +3,7 @@ module github.com/fly-examples/postgres-ha
 go 1.16
 
 require (
+	github.com/jackc/pgtype v1.6.2 // indirect
 	github.com/jackc/pgx/v4 v4.10.1
 	github.com/kr/pretty v0.2.0 // indirect
 	github.com/pkg/errors v0.8.1

pkg/check/check.go

@@ -0,0 +1,61 @@
+package check
+
+import (
+	"fmt"
+	"time"
+)
+
+type CheckFunction func() (string, error)
+type Check struct {
+	Name      string
+	CheckFunc CheckFunction
+	startTime time.Time
+	endTime   time.Time
+	message   string
+	err       error
+}
+
+func (c *Check) Process() {
+	c.startTime = time.Now()
+	c.message, c.err = c.CheckFunc()
+	c.endTime = time.Now()
+}
+
+func (c *Check) Error() string {
+	return c.err.Error()
+}
+
+func (c *Check) Passed() bool {
+	if c.startTime.IsZero() || c.endTime.IsZero() {
+		return false
+	}
+	return c.err == nil
+}
+
+func (c *Check) ExecutionTime() time.Duration {
+	if !c.endTime.IsZero() {
+		return RoundDuration(c.endTime.Sub(c.startTime), 2)
+	}
+	return RoundDuration(time.Now().Sub(c.startTime), 2)
+}
+
+func (c *Check) Result() string {
+	if c.startTime.IsZero() {
+		return fmt.Sprintf("[-] %s: %s", c.Name, "Not processed")
+	}
+	if !c.startTime.IsZero() && c.endTime.IsZero() {
+		return fmt.Sprintf("[✗] %s: %s (%s)", c.Name, "Timed out", c.ExecutionTime())
+	}
+	if c.Passed() {
+		return fmt.Sprintf("[✓] %s: %s (%s)", c.Name, c.message, c.ExecutionTime())
+	} else {
+		return fmt.Sprintf("[✗] %s: %s (%s)", c.Name, c.err.Error(), c.ExecutionTime())
+	}
+}
+
+func (c *Check) RawResult() string {
+	if c.Passed() {
+		return c.message
+	}
+	return c.err.Error()
+}

pkg/check/check_suite.go

@@ -0,0 +1,103 @@
+package check
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"time"
+)
+
+type OnCompletionHook func()
+type CheckSuite struct {
+	Name          string
+	Checks        []*Check
+	OnCompletion  OnCompletionHook
+	ErrOnSetup    error
+	executionTime time.Duration
+	processed     bool
+	clean         bool
+}
+
+func NewCheckSuite(name string) *CheckSuite {
+	return &CheckSuite{Name: name}
+}
+
+func (h *CheckSuite) Process(parentCtx context.Context) {
+	ctx, cancel := context.WithCancel(parentCtx)
+	start := time.Now()
+	for _, check := range h.Checks {
+		check.Process()
+	}
+	h.executionTime = RoundDuration(time.Since(start), 2)
+	h.processed = true
+	h.runOnCompletion()
+	cancel()
+
+	select {
+	case <-ctx.Done():
+		// Handle timeout
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			h.executionTime = RoundDuration(time.Since(start), 2)
+			h.processed = true
+			h.runOnCompletion()
+		}
+	}
+}
+
+func (h *CheckSuite) runOnCompletion() {
+	if h.clean {
+		return
+	}
+	if h.OnCompletion != nil {
+		h.OnCompletion()
+	}
+	h.clean = true
+}
+
+func (h *CheckSuite) AddCheck(name string, checkFunc CheckFunction) *Check {
+	check := &Check{Name: name, CheckFunc: checkFunc}
+	h.Checks = append(h.Checks, check)
+	return check
+}
+
+func (h *CheckSuite) Passed() bool {
+	for _, check := range h.Checks {
+		if !check.Passed() {
+			return false
+		}
+	}
+	return true
+}
+
+func (h *CheckSuite) Result() string {
+	checkStr := []string{}
+	for _, check := range h.Checks {
+		checkStr = append(checkStr, check.Result())
+	}
+	return strings.Join(checkStr, "\n")
+}
+
+func (h *CheckSuite) RawResult() string {
+	checkStr := []string{}
+	for _, check := range h.Checks {
+		checkStr = append(checkStr, check.RawResult())
+	}
+	return strings.Join(checkStr, "\n")
+}
+
+// Print will send output straight to stdout.
+func (h *CheckSuite) Print() {
+	if h.processed {
+		for _, check := range h.Checks {
+			fmt.Println(check.Result())
+		}
+		fmt.Printf("Total execution time of %q checks: %s\n", h.Name, h.executionTime)
+	} else {
+		if len(h.Checks) > 0 {
+			fmt.Printf("%q hasn't been processed. %d check(s) pending evaluation.\n", h.Name, len(h.Checks))
+		} else {
+			fmt.Printf("%q has no checks to evaluate.\n", h.Name)
+		}
+	}
+}