tracinginterceptor_test.go 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. package serverinterceptors
  2. import (
  3. "context"
  4. "sync"
  5. "sync/atomic"
  6. "testing"
  7. "github.com/stretchr/testify/assert"
  8. "github.com/tal-tech/go-zero/core/trace/tracespec"
  9. "google.golang.org/grpc"
  10. "google.golang.org/grpc/metadata"
  11. )
  12. func TestUnaryTracingInterceptor(t *testing.T) {
  13. interceptor := UnaryTracingInterceptor("foo")
  14. var run int32
  15. var wg sync.WaitGroup
  16. wg.Add(1)
  17. _, err := interceptor(context.Background(), nil, &grpc.UnaryServerInfo{
  18. FullMethod: "/",
  19. }, func(ctx context.Context, req interface{}) (interface{}, error) {
  20. defer wg.Done()
  21. atomic.AddInt32(&run, 1)
  22. return nil, nil
  23. })
  24. wg.Wait()
  25. assert.Nil(t, err)
  26. assert.Equal(t, int32(1), atomic.LoadInt32(&run))
  27. }
  28. func TestUnaryTracingInterceptor_GrpcFormat(t *testing.T) {
  29. interceptor := UnaryTracingInterceptor("foo")
  30. var wg sync.WaitGroup
  31. wg.Add(1)
  32. var md metadata.MD
  33. ctx := metadata.NewIncomingContext(context.Background(), md)
  34. _, err := interceptor(ctx, nil, &grpc.UnaryServerInfo{
  35. FullMethod: "/",
  36. }, func(ctx context.Context, req interface{}) (interface{}, error) {
  37. defer wg.Done()
  38. assert.True(t, len(ctx.Value(tracespec.TracingKey).(tracespec.Trace).TraceId()) > 0)
  39. assert.True(t, len(ctx.Value(tracespec.TracingKey).(tracespec.Trace).SpanId()) > 0)
  40. return nil, nil
  41. })
  42. wg.Wait()
  43. assert.Nil(t, err)
  44. }