github.com/lingyao2333/mo-zero@v1.4.1/rest/handler/tracinghandler_test.go (about)

     1  package handler
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"testing"
     8  
     9  	ztrace "github.com/lingyao2333/mo-zero/core/trace"
    10  	"github.com/lingyao2333/mo-zero/rest/chain"
    11  	"github.com/stretchr/testify/assert"
    12  	"go.opentelemetry.io/otel"
    13  	"go.opentelemetry.io/otel/propagation"
    14  	"go.opentelemetry.io/otel/trace"
    15  )
    16  
    17  func TestOtelHandler(t *testing.T) {
    18  	ztrace.StartAgent(ztrace.Config{
    19  		Name:     "go-zero-test",
    20  		Endpoint: "http://localhost:14268/api/traces",
    21  		Batcher:  "jaeger",
    22  		Sampler:  1.0,
    23  	})
    24  	defer ztrace.StopAgent()
    25  
    26  	for _, test := range []string{"", "bar"} {
    27  		t.Run(test, func(t *testing.T) {
    28  			h := chain.New(TracingHandler("foo", test)).Then(
    29  				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    30  					ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
    31  					spanCtx := trace.SpanContextFromContext(ctx)
    32  					assert.True(t, spanCtx.IsValid())
    33  				}))
    34  			ts := httptest.NewServer(h)
    35  			defer ts.Close()
    36  
    37  			client := ts.Client()
    38  			err := func(ctx context.Context) error {
    39  				ctx, span := otel.Tracer("httptrace/client").Start(ctx, "test")
    40  				defer span.End()
    41  
    42  				req, _ := http.NewRequest("GET", ts.URL, nil)
    43  				otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
    44  
    45  				res, err := client.Do(req)
    46  				assert.Nil(t, err)
    47  				return res.Body.Close()
    48  			}(context.Background())
    49  
    50  			assert.Nil(t, err)
    51  		})
    52  	}
    53  }
    54  
    55  func TestDontTracingSpanName(t *testing.T) {
    56  	ztrace.StartAgent(ztrace.Config{
    57  		Name:     "go-zero-test",
    58  		Endpoint: "http://localhost:14268/api/traces",
    59  		Batcher:  "jaeger",
    60  		Sampler:  1.0,
    61  	})
    62  
    63  	DontTraceSpan("bar")
    64  
    65  	for _, test := range []string{"", "bar", "foo"} {
    66  		t.Run(test, func(t *testing.T) {
    67  			h := chain.New(TracingHandler("foo", test)).Then(
    68  				http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    69  					spanCtx := trace.SpanContextFromContext(r.Context())
    70  					if test == "bar" {
    71  						assert.False(t, spanCtx.IsValid())
    72  						return
    73  					}
    74  
    75  					assert.True(t, spanCtx.IsValid())
    76  				}))
    77  			ts := httptest.NewServer(h)
    78  			defer ts.Close()
    79  
    80  			client := ts.Client()
    81  			err := func(ctx context.Context) error {
    82  				ctx, span := otel.Tracer("httptrace/client").Start(ctx, "test")
    83  				defer span.End()
    84  
    85  				req, _ := http.NewRequest("GET", ts.URL, nil)
    86  				otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))
    87  
    88  				res, err := client.Do(req)
    89  				assert.Nil(t, err)
    90  				return res.Body.Close()
    91  			}(context.Background())
    92  
    93  			assert.Nil(t, err)
    94  		})
    95  	}
    96  }