From 0cfb59fcdc65527b0e8f722c3dac58470b4ff9b7 Mon Sep 17 00:00:00 2001
From: "Riandy R.N" <riandyrn@users.noreply.github.com>
Date: Tue, 17 Sep 2024 12:42:03 +0300
Subject: [PATCH] test: control span context in trace response headers option
 to ensure the value of X-Trace-Sampled is correct;

---
 test/cases/sdk_test.go | 96 ++++++++++++++++++++++++++++--------------
 1 file changed, 65 insertions(+), 31 deletions(-)

diff --git a/test/cases/sdk_test.go b/test/cases/sdk_test.go
index f8ad07e..e116c78 100644
--- a/test/cases/sdk_test.go
+++ b/test/cases/sdk_test.go
@@ -2,6 +2,7 @@ package otelchi_test
 
 import (
 	"context"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -430,58 +431,91 @@ func TestSDKIntegrationWithDefaultLegacyHeaderKey(t *testing.T) {
 }
 
 func TestSDKIntegrationWithTraceResponseHeaders(t *testing.T) {
+	// prepare both sampled & non-sampled span context
+	spanCtxSampled := trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    [16]byte{1},
+		SpanID:     [8]byte{1},
+		Remote:     true,
+		TraceFlags: trace.FlagsSampled,
+	})
+	spanCtxNotSampled := trace.NewSpanContext(trace.SpanContextConfig{
+		TraceID:    [16]byte{2},
+		SpanID:     [8]byte{2},
+		Remote:     true,
+		TraceFlags: 0,
+	})
+
+	// define test cases
 	testCases := []struct {
-		Name                       string
-		TraceHeaderConfig          otelchi.TraceHeaderConfig
-		ExpTraceResponseIDKey      string
-		ExpTraceResponseSampledKey string
+		Name                          string
+		TraceHeaderConfig             otelchi.TraceHeaderConfig
+		SpanContext                   trace.SpanContext
+		ExpTraceResponseIDKey         string
+		ExpTraceResponseSampledKey    string
+		ExpTraceResponseSampledKeyVal bool
 	}{
 		{
-			Name:                       "Default Trace Config",
-			TraceHeaderConfig:          otelchi.TraceHeaderConfig{},
-			ExpTraceResponseIDKey:      otelchi.DefaultTraceIDResponseHeaderKey,
-			ExpTraceResponseSampledKey: otelchi.DefaultTraceSampledResponseHeaderKey,
+			Name:                          "Default Trace Config, Trace Sampled",
+			TraceHeaderConfig:             otelchi.TraceHeaderConfig{},
+			SpanContext:                   spanCtxSampled,
+			ExpTraceResponseIDKey:         otelchi.DefaultTraceIDResponseHeaderKey,
+			ExpTraceResponseSampledKey:    otelchi.DefaultTraceSampledResponseHeaderKey,
+			ExpTraceResponseSampledKeyVal: true,
+		},
+		{
+			Name:                          "Default Trace Config, Trace Not Sampled",
+			TraceHeaderConfig:             otelchi.TraceHeaderConfig{},
+			SpanContext:                   spanCtxNotSampled,
+			ExpTraceResponseIDKey:         otelchi.DefaultTraceIDResponseHeaderKey,
+			ExpTraceResponseSampledKey:    otelchi.DefaultTraceSampledResponseHeaderKey,
+			ExpTraceResponseSampledKeyVal: false,
+		},
+		{
+			Name: "Custom Trace Config, Trace Sampled",
+			TraceHeaderConfig: otelchi.TraceHeaderConfig{
+				TraceIDHeader:      "X-Custom-Trace-ID",
+				TraceSampledHeader: "X-Custom-Trace-Sampled",
+			},
+			SpanContext:                   spanCtxSampled,
+			ExpTraceResponseIDKey:         "X-Custom-Trace-ID",
+			ExpTraceResponseSampledKey:    "X-Custom-Trace-Sampled",
+			ExpTraceResponseSampledKeyVal: true,
 		},
 		{
-			Name: "Custom Trace Config",
+			Name: "Custom Trace Config, Trace Not Sampled",
 			TraceHeaderConfig: otelchi.TraceHeaderConfig{
 				TraceIDHeader:      "X-Custom-Trace-ID",
 				TraceSampledHeader: "X-Custom-Trace-Sampled",
 			},
-			ExpTraceResponseIDKey:      "X-Custom-Trace-ID",
-			ExpTraceResponseSampledKey: "X-Custom-Trace-Sampled",
+			SpanContext:                   spanCtxNotSampled,
+			ExpTraceResponseIDKey:         "X-Custom-Trace-ID",
+			ExpTraceResponseSampledKey:    "X-Custom-Trace-Sampled",
+			ExpTraceResponseSampledKeyVal: false,
 		},
 	}
 	for _, testCase := range testCases {
 		t.Run(testCase.Name, func(t *testing.T) {
-			router, sr := newSDKTestRouter(
-				"foobar",
-				true,
-				otelchi.WithTraceResponseHeaders(testCase.TraceHeaderConfig),
+			// configure router
+			router := chi.NewRouter()
+			router.Use(
+				otelchi.Middleware(
+					"foobar",
+					otelchi.WithChiRoutes(router),
+					otelchi.WithTraceResponseHeaders(testCase.TraceHeaderConfig),
+				),
 			)
 			router.HandleFunc("/user/{id:[0-9]+}", ok)
 			router.HandleFunc("/book/{title}", ok)
 
+			// execute requests
 			r0 := httptest.NewRequest("GET", "/user/123", nil)
+			r0 = r0.WithContext(trace.ContextWithRemoteSpanContext(context.Background(), testCase.SpanContext))
 			w := httptest.NewRecorder()
 			router.ServeHTTP(w, r0)
 
-			recordedSpans := sr.Ended()
-			require.Len(t, recordedSpans, 1)
-			checkSpans(t, recordedSpans, []spanValueCheck{
-				{
-					Name: "/user/{id:[0-9]+}",
-					Kind: trace.SpanKindServer,
-					Attributes: getSemanticAttributes(
-						"foobar",
-						http.StatusOK,
-						"GET",
-						"/user/{id:[0-9]+}",
-					),
-				},
-			})
-			require.Equal(t, recordedSpans[0].SpanContext().TraceID().String(), w.Header().Get(testCase.ExpTraceResponseIDKey))
-			require.Equal(t, "true", w.Header().Get(testCase.ExpTraceResponseSampledKey))
+			// check response headers
+			require.Equal(t, testCase.SpanContext.TraceID().String(), w.Header().Get(testCase.ExpTraceResponseIDKey))
+			require.Equal(t, fmt.Sprintf("%v", testCase.ExpTraceResponseSampledKeyVal), w.Header().Get(testCase.ExpTraceResponseSampledKey))
 		})
 	}
 }