github.com/gogf/gf/v2@v2.7.4/net/gclient/gclient_z_unit_feature_trace_test.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/gogf/gf.
     6  
     7  package gclient_test
     8  
     9  import (
    10  	"context"
    11  	"fmt"
    12  	"net/http"
    13  	"testing"
    14  	"time"
    15  
    16  	"go.opentelemetry.io/otel"
    17  	sdkTrace "go.opentelemetry.io/otel/sdk/trace"
    18  	"go.opentelemetry.io/otel/trace"
    19  
    20  	"github.com/gogf/gf/v2/frame/g"
    21  	"github.com/gogf/gf/v2/internal/tracing"
    22  	"github.com/gogf/gf/v2/net/ghttp"
    23  	"github.com/gogf/gf/v2/test/gtest"
    24  	"github.com/gogf/gf/v2/util/guid"
    25  )
    26  
    27  type CustomProvider struct {
    28  	*sdkTrace.TracerProvider
    29  }
    30  
    31  func NewCustomProvider() *CustomProvider {
    32  	return &CustomProvider{
    33  		TracerProvider: sdkTrace.NewTracerProvider(
    34  			sdkTrace.WithIDGenerator(NewCustomIDGenerator()),
    35  		),
    36  	}
    37  }
    38  
    39  type CustomIDGenerator struct{}
    40  
    41  func NewCustomIDGenerator() *CustomIDGenerator {
    42  	return &CustomIDGenerator{}
    43  }
    44  
    45  func (id *CustomIDGenerator) NewIDs(ctx context.Context) (traceID trace.TraceID, spanID trace.SpanID) {
    46  	return tracing.NewIDs()
    47  }
    48  
    49  func (id *CustomIDGenerator) NewSpanID(ctx context.Context, traceID trace.TraceID) (spanID trace.SpanID) {
    50  	return tracing.NewSpanID()
    51  }
    52  
    53  func TestClient_CustomProvider(t *testing.T) {
    54  	provider := otel.GetTracerProvider()
    55  	defer otel.SetTracerProvider(provider)
    56  
    57  	otel.SetTracerProvider(NewCustomProvider())
    58  
    59  	s := g.Server(guid.S())
    60  	s.BindHandler("/hello", func(r *ghttp.Request) {
    61  		r.Response.WriteHeader(200)
    62  		r.Response.WriteJson(g.Map{"field": "test_for_response_body"})
    63  	})
    64  	s.SetDumpRouterMap(false)
    65  	s.Start()
    66  	defer s.Shutdown()
    67  
    68  	time.Sleep(100 * time.Millisecond)
    69  	gtest.C(t, func(t *gtest.T) {
    70  		c := g.Client()
    71  		url := fmt.Sprintf("127.0.0.1:%d/hello", s.GetListenedPort())
    72  		resp, err := c.DoRequest(ctx, http.MethodGet, url)
    73  		t.AssertNil(err)
    74  		t.AssertNE(resp, nil)
    75  		t.Assert(resp.ReadAllString(), "{\"field\":\"test_for_response_body\"}")
    76  	})
    77  }