更新

middleware/traceId_test.go

@@ -0,0 +1,154 @@
+package middleware_test
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/google/uuid"
+	"github.com/yuninks/loggerx"
+	"github.com/yuninks/loggerx/middleware"
+)
+
+func newTestEngine(log *loggerx.Logger) *gin.Engine {
+	gin.SetMode(gin.TestMode)
+	g := gin.New()
+	g.Use(middleware.SetGinTraceId(log))
+	return g
+}
+
+// trace id 应当写回响应头，并在 handler 里可读
+func TestGinTraceIdSetsResponseHeader(t *testing.T) {
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(t.TempDir()))
+	defer log.Close()
+
+	var seen string
+	g := newTestEngine(log)
+	g.GET("/ping", func(c *gin.Context) {
+		seen = middleware.GetTraceId(c.Request.Context(), log.GetTraceField())
+		c.Status(200)
+	})
+
+	w := httptest.NewRecorder()
+	g.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ping", nil))
+
+	header := w.Header().Get(middleware.TraceHeader)
+	if header == "" {
+		t.Fatalf("响应头 %s 没有写回，跨服务排障拿不到同一个 id", middleware.TraceHeader)
+	}
+	if _, err := uuid.Parse(header); err != nil {
+		t.Errorf("生成的 trace id 不是合法 uuid: %q", header)
+	}
+	if seen != header {
+		t.Errorf("handler 里读到的 trace id (%q) 与响应头 (%q) 不一致", seen, header)
+	}
+}
+
+// 上游带 trace id 时必须透传，不能重新生成
+func TestGinTraceIdPropagatesFromHeader(t *testing.T) {
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(t.TempDir()))
+	defer log.Close()
+
+	const upstream = "8f3d1c22-1111-4222-8333-444455556666"
+	var seen string
+	g := newTestEngine(log)
+	g.GET("/ping", func(c *gin.Context) {
+		seen = middleware.GetTraceId(c.Request.Context(), log.GetTraceField())
+		c.Status(200)
+	})
+
+	req := httptest.NewRequest(http.MethodGet, "/ping", nil)
+	req.Header.Set(middleware.TraceHeader, upstream)
+	w := httptest.NewRecorder()
+	g.ServeHTTP(w, req)
+
+	if seen != upstream {
+		t.Errorf("上游 trace id 未透传: 期望 %q 实际 %q", upstream, seen)
+	}
+	if got := w.Header().Get(middleware.TraceHeader); got != upstream {
+		t.Errorf("响应头应回传上游 trace id，实际 %q", got)
+	}
+}
+
+// 日志里应当带上 middleware 注入的 trace id（gin.Context 与 request ctx 两条路径）
+func TestTraceIdAppearsInLog(t *testing.T) {
+	dir := t.TempDir()
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(dir))
+	defer log.Close()
+
+	g := newTestEngine(log)
+	g.GET("/ping", func(c *gin.Context) {
+		log.Info(c, "via-gin-ctx")
+		log.Info(c.Request.Context(), "via-request-ctx")
+		c.Status(200)
+	})
+
+	w := httptest.NewRecorder()
+	g.ServeHTTP(w, httptest.NewRequest(http.MethodGet, "/ping", nil))
+	traceId := w.Header().Get(middleware.TraceHeader)
+	if traceId == "" {
+		t.Fatal("没有 trace id")
+	}
+	if err := log.MustSync(); err != nil {
+		t.Fatal(err)
+	}
+
+	content := readAll(t, dir)
+	if !contains(content, traceId) {
+		t.Errorf("日志里没有出现 trace id %q:\n%s", traceId, content)
+	}
+	if !contains(content, "via-gin-ctx") || !contains(content, "via-request-ctx") {
+		t.Errorf("两条日志都应当写入:\n%s", content)
+	}
+}
+
+// SetTraceId 打标：已存在时不覆盖，取不到时返回空串而不是 panic
+func TestSetTraceIdByKeySemantics(t *testing.T) {
+	log := loggerx.NewLogger(context.Background(), loggerx.SetDir(t.TempDir()))
+	defer log.Close()
+
+	ctx := middleware.SetTraceId(context.Background(), log)
+	first := middleware.GetTraceId(ctx, log.GetTraceField())
+	if first == "" {
+		t.Fatal("SetTraceId 应当生成一个 id")
+	}
+
+	again := middleware.SetTraceId(ctx, log)
+	if got := middleware.GetTraceId(again, log.GetTraceField()); got != first {
+		t.Errorf("已有 trace id 时不该覆盖: %q -> %q", first, got)
+	}
+
+	if got := middleware.GetTraceId(nil, "trace_id"); got != "" {
+		t.Errorf("nil ctx 应返回空串，实际 %q", got)
+	}
+	if got := middleware.GetTraceId(context.Background(), "trace_id"); got != "" {
+		t.Errorf("没有值时该返回空串，实际 %q", got)
+	}
+}
+
+func readAll(t *testing.T, dir string) string {
+	t.Helper()
+	var sb []byte
+	files, _ := filepath.Glob(filepath.Join(dir, "*.log"))
+	for _, f := range files {
+		b, err := os.ReadFile(f)
+		if err != nil {
+			t.Fatalf("读取 %s: %v", f, err)
+		}
+		sb = append(sb, b...)
+	}
+	return string(sb)
+}
+
+func contains(s, sub string) bool {
+	for i := 0; i+len(sub) <= len(s); i++ {
+		if s[i:i+len(sub)] == sub {
+			return true
+		}
+	}
+	return false
+}