go-omni/health/health_test.go

package health

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"net/http/httptest"
	"testing"
	"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"
)

func TestCheckerReadinessReportsAllProbes(t *testing.T) {
	checker, err := New(
		Probe{Name: "database", Check: func(context.Context) error { return nil }},
		Probe{Name: "cache", Check: func(context.Context) error { return errors.New("connection refused") }},
	)
	if err != nil {
		t.Fatal(err)
	}
	app := omni.NewWithRouter(nethttp.New())
	app.GET("/live", checker.LivenessHandler())
	app.GET("/ready", checker.ReadinessHandler())

	live := httptest.NewRecorder()
	app.ServeHTTP(live, httptest.NewRequest(http.MethodGet, "/live", nil))
	if live.Code != http.StatusOK {
		t.Fatalf("expected liveness 200, got %d", live.Code)
	}

	ready := httptest.NewRecorder()
	app.ServeHTTP(ready, httptest.NewRequest(http.MethodGet, "/ready", nil))
	if ready.Code != http.StatusServiceUnavailable {
		t.Fatalf("expected readiness 503, got %d", ready.Code)
	}
	var body struct {
		Code string `json:"code"`
		Data Result `json:"data"`
	}
	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 != "" {
		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",
		Timeout: 5 * time.Millisecond,
		Check: func(ctx context.Context) error {
			<-ctx.Done()
			return ctx.Err()
		},
	})
	if err != nil {
		t.Fatal(err)
	}
	result := checker.Check(context.Background())
	if result.Ready() || result.Checks[0].Status != StatusNotReady {
		t.Fatalf("expected timeout to fail readiness: %+v", result)
	}

	for _, probe := range []Probe{{Check: func(context.Context) error { return nil }}, {Name: "slow", Check: func(context.Context) error { return nil }}, {Name: "negative", Timeout: -time.Second, Check: func(context.Context) error { return nil }}} {
		if err := checker.Add(probe); err == nil {
			t.Fatalf("expected registration error for %+v", probe)
		}
	}
}