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) } } }