loggerx/middleware/traceId_test.go

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
}