tracinginterceptor_test.go 1.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. package clientinterceptors
  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"
  9. "google.golang.org/grpc"
  10. "google.golang.org/grpc/metadata"
  11. )
  12. func TestTracingInterceptor(t *testing.T) {
  13. var run int32
  14. var wg sync.WaitGroup
  15. wg.Add(1)
  16. cc := new(grpc.ClientConn)
  17. err := TracingInterceptor(context.Background(), "/foo", nil, nil, cc,
  18. func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
  19. opts ...grpc.CallOption) error {
  20. defer wg.Done()
  21. atomic.AddInt32(&run, 1)
  22. return nil
  23. })
  24. wg.Wait()
  25. assert.Nil(t, err)
  26. assert.Equal(t, int32(1), atomic.LoadInt32(&run))
  27. }
  28. func TestTracingInterceptor_GrpcFormat(t *testing.T) {
  29. var run int32
  30. var wg sync.WaitGroup
  31. wg.Add(1)
  32. md := metadata.New(map[string]string{
  33. "foo": "bar",
  34. })
  35. carrier, err := trace.Inject(trace.GrpcFormat, md)
  36. assert.Nil(t, err)
  37. ctx, _ := trace.StartServerSpan(context.Background(), carrier, "user", "/foo")
  38. cc := new(grpc.ClientConn)
  39. err = TracingInterceptor(ctx, "/foo", nil, nil, cc,
  40. func(ctx context.Context, method string, req, reply interface{}, cc *grpc.ClientConn,
  41. opts ...grpc.CallOption) error {
  42. defer wg.Done()
  43. atomic.AddInt32(&run, 1)
  44. return nil
  45. })
  46. wg.Wait()
  47. assert.Nil(t, err)
  48. assert.Equal(t, int32(1), atomic.LoadInt32(&run))
  49. }