github.com/mier85/go-sensor@v1.30.1-0.20220920111756-9bf41b3bc7e0/w3ctrace/middleware_test.go (about)

     1  // (c) Copyright IBM Corp. 2021
     2  // (c) Copyright Instana Inc. 2020
     3  
     4  package w3ctrace_test
     5  
     6  import (
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"testing"
    10  
    11  	"github.com/instana/testify/assert"
    12  	"github.com/instana/testify/require"
    13  	"github.com/mier85/go-sensor/w3ctrace"
    14  )
    15  
    16  func TestTracingHandlerFunc(t *testing.T) {
    17  	h := w3ctrace.TracingHandlerFunc(func(w http.ResponseWriter, req *http.Request) {
    18  		w.Header().Set("X-Test", "value")
    19  		w.Write([]byte("Ok"))
    20  	})
    21  
    22  	resp := httptest.NewRecorder()
    23  
    24  	req := httptest.NewRequest("GET", "/", nil)
    25  	req.Header.Set("traceparent", exampleTraceParent)
    26  	req.Header.Set("tracestate", exampleTraceState)
    27  
    28  	h(resp, req)
    29  	require.Equal(t, http.StatusOK, resp.Result().StatusCode)
    30  	assert.Equal(t, "Ok", resp.Body.String())
    31  
    32  	assert.Equal(t, "value", resp.Header().Get("X-Test"))
    33  	assert.Equal(t, exampleTraceParent, resp.Header().Get(w3ctrace.TraceParentHeader))
    34  	assert.Equal(t, exampleTraceState, resp.Header().Get(w3ctrace.TraceStateHeader))
    35  }
    36  
    37  func TestTracingHandlerFunc_NoContext(t *testing.T) {
    38  	h := w3ctrace.TracingHandlerFunc(func(w http.ResponseWriter, req *http.Request) {
    39  		w.Header().Set("X-Test", "value")
    40  		w.Write([]byte("Ok"))
    41  	})
    42  
    43  	resp := httptest.NewRecorder()
    44  
    45  	req := httptest.NewRequest("GET", "/", nil)
    46  
    47  	h(resp, req)
    48  	require.Equal(t, http.StatusOK, resp.Result().StatusCode)
    49  	assert.Equal(t, "Ok", resp.Body.String())
    50  
    51  	assert.Equal(t, "value", resp.Header().Get("X-Test"))
    52  	assert.Empty(t, resp.Header().Get(w3ctrace.TraceParentHeader))
    53  	assert.Empty(t, resp.Header().Get(w3ctrace.TraceStateHeader))
    54  }