From fd782b25c3da95ba14154b883f7a02a8d0d98d18 Mon Sep 17 00:00:00 2001
From: yangyou <tjbuggo@gmail.com>
Date: Tue, 20 Aug 2024 22:42:04 +0800
Subject: [PATCH] Add test for reading trace from http request header

---
 logtracing/span.go    |  4 ++++
 server/logger_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/logtracing/span.go b/logtracing/span.go
index 7626df8..2646245 100644
--- a/logtracing/span.go
+++ b/logtracing/span.go
@@ -74,6 +74,10 @@ func (s *span) AppendKVs(keyvals ...interface{}) {
 	s.keyvals = append(s.keyvals, keyvals...)
 }
 
+func (s *span) TraceID() string {
+	return s.traceID.String()
+}
+
 func (s *span) meta() spanMeta {
 	return spanMeta{
 		TraceID:   s.traceID,
diff --git a/server/logger_test.go b/server/logger_test.go
index 4eaca38..a2f59b4 100644
--- a/server/logger_test.go
+++ b/server/logger_test.go
@@ -1,12 +1,16 @@
 package server
 
 import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
 	"github.com/theplant/appkit/contexts"
 	"github.com/theplant/appkit/log"
+	"github.com/theplant/appkit/logtracing"
 )
 
 func TestLogRequest(t *testing.T) {
@@ -32,4 +36,29 @@ func TestLogRequest(t *testing.T) {
 	if rw.Result().StatusCode != http.StatusOK {
 		t.Fatalf("unexpected status code: %d", rw.Result().StatusCode)
 	}
+
+	r16 := make([]byte, 16)
+	r8 := make([]byte, 8)
+	rand.Read(r16)
+	rand.Read(r8)
+	traceID := hex.EncodeToString(r16)
+	spanID := hex.EncodeToString(r8)
+	trace := fmt.Sprintf("%s-%s-%s-%s", "00", traceID, spanID, "00")
+	req, err = http.NewRequest("GET", "http://example.com", nil)
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+	req.Header.Add(traceHeaderKey, trace)
+	h = Compose(
+		// Recovery should come before logReq to set the status code to 500
+		Recovery,
+		LogRequest,
+	)(http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		_, span := logtracing.StartSpan(r.Context(), "")
+		if span.TraceID() != traceID {
+			t.Errorf("traceID should be: %s, but got: %s", traceID, span.TraceID())
+		}
+	}))
+
+	h.ServeHTTP(rw, req)
 }