feat: add dependency shutdown and configurable service endpoints

health/health_test.go

@@ -1,6 +1,7 @@
 package health
 
 import (
+	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -10,6 +11,9 @@ import (
 	"time"
 
 	omni "git.yuntuo.hk/k2safe/go-omni"
+	apierrors "git.yuntuo.hk/k2safe/go-omni/errors"
+	"git.yuntuo.hk/k2safe/go-omni/httpx"
+	"git.yuntuo.hk/k2safe/go-omni/response"
 	"git.yuntuo.hk/k2safe/go-omni/transport/nethttp"
 )
 
@@ -43,11 +47,119 @@ func TestCheckerReadinessReportsAllProbes(t *testing.T) {
 	if err := json.Unmarshal(ready.Body.Bytes(), &body); err != nil {
 		t.Fatal(err)
 	}
-	if body.Code != "NOT_READY" || body.Data.Status != StatusNotReady || len(body.Data.Checks) != 2 || body.Data.Checks[1].Error != "connection refused" {
+	if body.Code != "NOT_READY" || body.Data.Status != StatusNotReady || len(body.Data.Checks) != 2 || body.Data.Checks[1].Error != "" {
 		t.Fatalf("unexpected readiness response: %+v", body)
 	}
 }
 
+func TestProbeResponsePolicies(t *testing.T) {
+	for _, raw := range []bool{false, true} {
+		t.Run(map[bool]string{true: "raw", false: "app"}[raw], func(t *testing.T) {
+			probeErr := errors.New("private database connection details")
+			checker, _ := New(Probe{Name: "database", Check: func(context.Context) error { return probeErr }})
+			app := omni.New(omni.WithResponse(response.Config{Body: map[string]any{"code": "${code}", "msg": "${message}", "data": "${data}", "httpStatus": "${status}"}, CodeType: "int32", SuccessCode: 200, ErrorCode: 9999, SuccessMessage: "操作成功", DataMode: "always"}), omni.WithErrors(apierrors.Catalog{DefaultStatus: 200, Definitions: map[string]apierrors.Definition{"NOT_READY": {Code: 503001, Message: "暂未就绪"}}}))
+			var failure *apierrors.APIError
+			app.Use(func(next httpx.Handler) httpx.Handler {
+				return func(c *httpx.Context) error { err := next(c); failure = c.Failure(); return err }
+			})
+			var opts []HandlerOption
+			if raw {
+				opts = append(opts, WithRawResponse())
+			}
+			app.GET("/ready", checker.ReadinessHandler(opts...))
+			app.GET("/live", checker.LivenessHandler(opts...))
+			for _, path := range []string{"/ready", "/live"} {
+				out := httptest.NewRecorder()
+				app.ServeHTTP(out, httptest.NewRequest("GET", path, nil))
+				want := 200
+				if path == "/ready" {
+					want = 503
+				}
+				if out.Code != want {
+					t.Fatal(out.Code, out.Body.String())
+				}
+				var body map[string]any
+				if err := json.Unmarshal(out.Body.Bytes(), &body); err != nil {
+					t.Fatal(err)
+				}
+				if raw {
+					if _, ok := body["code"]; ok || body["status"] == nil {
+						t.Fatal(body)
+					}
+				} else {
+					if body["httpStatus"] != float64(want) {
+						t.Fatal(body)
+					}
+					if want == 503 && (body["code"] != float64(503001) || body["msg"] != "暂未就绪") {
+						t.Fatal(body)
+					}
+					if want == 200 && (body["code"] != float64(200) || body["msg"] != "操作成功") {
+						t.Fatal(body)
+					}
+				}
+				if want == 503 && (failure == nil || failure.Status != 503 || !errors.Is(failure, probeErr)) {
+					t.Fatal("failure/log cause lost", failure)
+				}
+				if bytes.Contains(out.Body.Bytes(), []byte(probeErr.Error())) {
+					t.Fatal("probe cause leaked")
+				}
+			}
+		})
+	}
+}
+
+func TestShutdownReadinessAndCustomEncoderStatus(t *testing.T) {
+	ctx, stop := context.WithCancel(context.Background())
+	defer stop()
+	checker, _ := New()
+	app := omni.New(omni.WithResponseEncoder(response.EncoderFunc(func(r response.Result) (response.Payload, error) {
+		return response.JSON(200, map[string]any{"status": r.Status, "data": r.Data})
+	})))
+	app.GET("/ready", checker.ReadinessHandler(WithShutdownContext(ctx)))
+	app.GET("/live", checker.LivenessHandler())
+	stop()
+	out := httptest.NewRecorder()
+	app.ServeHTTP(out, httptest.NewRequest("GET", "/ready", nil))
+	if out.Code != 503 {
+		t.Fatal("encoder masked failed readiness", out.Code)
+	}
+	out = httptest.NewRecorder()
+	app.ServeHTTP(out, httptest.NewRequest("GET", "/live", nil))
+	if out.Code != 200 {
+		t.Fatal("shutdown changed liveness", out.Code)
+	}
+}
+
+func TestBrokenProbeEncoderCannotReportHealthy(t *testing.T) {
+	for _, panicEncoder := range []bool{false, true} {
+		for _, unhealthy := range []bool{false, true} {
+			checker, _ := New()
+			if unhealthy {
+				_ = checker.Add(Probe{Name: "dependency", Check: func(context.Context) error { return errors.New("offline") }})
+			}
+			app := omni.New(omni.WithResponseEncoder(response.EncoderFunc(func(r response.Result) (response.Payload, error) {
+				if r.Data != nil {
+					if panicEncoder {
+						panic("private encoder detail")
+					}
+					return response.Payload{}, errors.New("private encoder detail")
+				}
+				return response.JSON(200, map[string]bool{"acknowledged": true})
+			})))
+			app.GET("/ready", checker.ReadinessHandler())
+			out := httptest.NewRecorder()
+			app.ServeHTTP(out, httptest.NewRequest("GET", "/ready", nil))
+			want := 500
+			if unhealthy {
+				want = 503
+			}
+			if out.Code != want || out.Body.Len() != 0 {
+				t.Fatal(out.Code, out.Body.String())
+			}
+		}
+	}
+}
+
 func TestCheckerTimeoutAndRegistrationValidation(t *testing.T) {
 	checker, err := New(Probe{
 		Name:    "slow",