feat(trace): add trace test helpers (#3108)

core/stores/sqlx/sqlconn_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/DATA-DOG/go-sqlmock"
 	"github.com/stretchr/testify/assert"
 	"github.com/zeromicro/go-zero/core/logx"
+	"github.com/zeromicro/go-zero/core/trace/tracetest"
 )
 
 const mockedDatasource = "sqlmock"
@@ -17,6 +18,7 @@ func init() {
 }
 
 func TestSqlConn(t *testing.T) {
+	me := tracetest.NewInMemoryExporter(t)
 	mock, err := buildConn()
 	assert.Nil(t, err)
 	mock.ExpectExec("any")
@@ -49,6 +51,7 @@ func TestSqlConn(t *testing.T) {
 	assert.NotNil(t, badConn.Transact(func(session Session) error {
 		return nil
 	}))
+	assert.Equal(t, 14, len(me.GetSpans()))
 }
 
 func buildConn() (mock sqlmock.Sqlmock, err error) {

core/trace/tracetest/tracetest.go

@@ -0,0 +1,20 @@
+package tracetest
+
+import (
+	"testing"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/sdk/trace"
+	"go.opentelemetry.io/otel/sdk/trace/tracetest"
+)
+
+// NewInMemoryExporter returns a new InMemoryExporter
+// and sets it as the global for tests.
+func NewInMemoryExporter(t *testing.T) *tracetest.InMemoryExporter {
+	me := tracetest.NewInMemoryExporter()
+	t.Cleanup(func() {
+		me.Reset()
+	})
+	otel.SetTracerProvider(trace.NewTracerProvider(trace.WithSyncer(me)))
+	return me
+}

rest/handler/tracehandler_test.go

@@ -10,9 +10,12 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	ztrace "github.com/zeromicro/go-zero/core/trace"
+	"github.com/zeromicro/go-zero/core/trace/tracetest"
 	"github.com/zeromicro/go-zero/rest/chain"
 	"go.opentelemetry.io/otel"
+	tcodes "go.opentelemetry.io/otel/codes"
 	"go.opentelemetry.io/otel/propagation"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -54,6 +57,31 @@ func TestOtelHandler(t *testing.T) {
 	}
 }
 
+func TestTraceHandler(t *testing.T) {
+	me := tracetest.NewInMemoryExporter(t)
+	h := chain.New(TraceHandler("foo", "/")).Then(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	ts := httptest.NewServer(h)
+	defer ts.Close()
+
+	client := ts.Client()
+	err := func(ctx context.Context) error {
+		req, _ := http.NewRequest("GET", ts.URL, nil)
+
+		res, err := client.Do(req)
+		assert.Nil(t, err)
+		return res.Body.Close()
+	}(context.Background())
+	assert.NoError(t, err)
+	assert.Equal(t, 1, len(me.GetSpans()))
+	span := me.GetSpans()[0].Snapshot()
+	assert.Equal(t, sdktrace.Status{
+		Code: tcodes.Unset,
+	}, span.Status())
+	assert.Equal(t, 0, len(span.Events()))
+	assert.Equal(t, 9, len(span.Attributes()))
+}
+
 func TestDontTracingSpan(t *testing.T) {
 	ztrace.StartAgent(ztrace.Config{
 		Name:     "go-zero-test",

rest/httpc/requests_test.go

@@ -10,9 +10,12 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	ztrace "github.com/zeromicro/go-zero/core/trace"
+	"github.com/zeromicro/go-zero/core/trace/tracetest"
 	"github.com/zeromicro/go-zero/rest/httpx"
 	"github.com/zeromicro/go-zero/rest/internal/header"
 	"github.com/zeromicro/go-zero/rest/router"
+	tcodes "go.opentelemetry.io/otel/codes"
+	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/trace"
 )
 
@@ -59,6 +62,7 @@ func TestDoRequest_Moved(t *testing.T) {
 }
 
 func TestDo(t *testing.T) {
+	me := tracetest.NewInMemoryExporter(t)
 	type Data struct {
 		Key    string `path:"key"`
 		Value  int    `form:"value"`
@@ -86,6 +90,13 @@ func TestDo(t *testing.T) {
 	resp, err := Do(context.Background(), http.MethodPost, svr.URL+"/nodes/:key", data)
 	assert.Nil(t, err)
 	assert.Equal(t, http.StatusOK, resp.StatusCode)
+	assert.Equal(t, 1, len(me.GetSpans()))
+	span := me.GetSpans()[0].Snapshot()
+	assert.Equal(t, sdktrace.Status{
+		Code: tcodes.Unset,
+	}, span.Status())
+	assert.Equal(t, 0, len(span.Events()))
+	assert.Equal(t, 7, len(span.Attributes()))
 }
 
 func TestDo_Ptr(t *testing.T) {

zrpc/internal/clientinterceptors/tracinginterceptor_test.go

@@ -9,7 +9,12 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/trace"
+	ztrace "github.com/zeromicro/go-zero/core/trace"
+	"github.com/zeromicro/go-zero/core/trace/tracetest"
+	"go.opentelemetry.io/otel/attribute"
+	tcodes "go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/sdk/trace"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -17,13 +22,13 @@ import (
 )
 
 func TestOpenTracingInterceptor(t *testing.T) {
-	trace.StartAgent(trace.Config{
+	ztrace.StartAgent(ztrace.Config{
 		Name:     "go-zero-test",
 		Endpoint: "http://localhost:14268/api/traces",
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
-	defer trace.StopAgent()
+	defer ztrace.StopAgent()
 
 	cc := new(grpc.ClientConn)
 	ctx := metadata.NewOutgoingContext(context.Background(), metadata.MD{})
@@ -36,20 +41,79 @@ func TestOpenTracingInterceptor(t *testing.T) {
 }
 
 func TestUnaryTracingInterceptor(t *testing.T) {
-	var run int32
-	var wg sync.WaitGroup
-	wg.Add(1)
-	cc := new(grpc.ClientConn)
-	err := UnaryTracingInterceptor(context.Background(), "/foo", nil, nil, cc,
-		func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
-			opts ...grpc.CallOption) error {
-			defer wg.Done()
-			atomic.AddInt32(&run, 1)
-			return nil
-		})
-	wg.Wait()
-	assert.Nil(t, err)
-	assert.Equal(t, int32(1), atomic.LoadInt32(&run))
+	t.Run("normal", func(t *testing.T) {
+		var run int32
+		cc := new(grpc.ClientConn)
+		me := tracetest.NewInMemoryExporter(t)
+		err := UnaryTracingInterceptor(context.Background(), "/proto.Hello/Echo",
+			nil, nil, cc,
+			func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
+				opts ...grpc.CallOption) error {
+				atomic.AddInt32(&run, 1)
+				return nil
+			})
+		assert.Nil(t, err)
+		assert.Equal(t, int32(1), atomic.LoadInt32(&run))
+
+		assert.Equal(t, 1, len(me.GetSpans()))
+		span := me.GetSpans()[0].Snapshot()
+		assert.Equal(t, 2, len(span.Events()))
+		assert.ElementsMatch(t, []attribute.KeyValue{
+			ztrace.RPCSystemGRPC,
+			semconv.RPCServiceKey.String("proto.Hello"),
+			semconv.RPCMethodKey.String("Echo"),
+			ztrace.StatusCodeAttr(codes.OK),
+		}, span.Attributes())
+	})
+
+	t.Run("grpc error status", func(t *testing.T) {
+		me := tracetest.NewInMemoryExporter(t)
+		cc := new(grpc.ClientConn)
+		err := UnaryTracingInterceptor(context.Background(), "/proto.Hello/Echo",
+			nil, nil, cc,
+			func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
+				opts ...grpc.CallOption) error {
+				return status.Error(codes.Unknown, "test")
+			})
+		assert.Error(t, err)
+		assert.Equal(t, 1, len(me.GetSpans()))
+		span := me.GetSpans()[0].Snapshot()
+		assert.Equal(t, trace.Status{
+			Code:        tcodes.Error,
+			Description: "test",
+		}, span.Status())
+		assert.Equal(t, 2, len(span.Events()))
+		assert.ElementsMatch(t, []attribute.KeyValue{
+			ztrace.RPCSystemGRPC,
+			semconv.RPCServiceKey.String("proto.Hello"),
+			semconv.RPCMethodKey.String("Echo"),
+			ztrace.StatusCodeAttr(codes.Unknown),
+		}, span.Attributes())
+	})
+
+	t.Run("non grpc status error", func(t *testing.T) {
+		me := tracetest.NewInMemoryExporter(t)
+		cc := new(grpc.ClientConn)
+		err := UnaryTracingInterceptor(context.Background(), "/proto.Hello/Echo",
+			nil, nil, cc,
+			func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn,
+				opts ...grpc.CallOption) error {
+				return errors.New("test")
+			})
+		assert.Error(t, err)
+		assert.Equal(t, 1, len(me.GetSpans()))
+		span := me.GetSpans()[0].Snapshot()
+		assert.Equal(t, trace.Status{
+			Code:        tcodes.Error,
+			Description: "test",
+		}, span.Status())
+		assert.Equal(t, 2, len(span.Events()))
+		assert.ElementsMatch(t, []attribute.KeyValue{
+			ztrace.RPCSystemGRPC,
+			semconv.RPCServiceKey.String("proto.Hello"),
+			semconv.RPCMethodKey.String("Echo"),
+		}, span.Attributes())
+	})
 }
 
 func TestUnaryTracingInterceptor_WithError(t *testing.T) {

zrpc/internal/serverinterceptors/tracinginterceptor_test.go

@@ -9,7 +9,12 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
-	"github.com/zeromicro/go-zero/core/trace"
+	ztrace "github.com/zeromicro/go-zero/core/trace"
+	"github.com/zeromicro/go-zero/core/trace/tracetest"
+	"go.opentelemetry.io/otel/attribute"
+	tcodes "go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/sdk/trace"
+	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -26,13 +31,13 @@ func TestUnaryOpenTracingInterceptor_Disable(t *testing.T) {
 }
 
 func TestUnaryOpenTracingInterceptor_Enabled(t *testing.T) {
-	trace.StartAgent(trace.Config{
+	ztrace.StartAgent(ztrace.Config{
 		Name:     "go-zero-test",
 		Endpoint: "http://localhost:14268/api/traces",
 		Batcher:  "jaeger",
 		Sampler:  1.0,
 	})
-	defer trace.StopAgent()
+	defer ztrace.StopAgent()
 
 	_, err := UnaryTracingInterceptor(context.Background(), nil, &grpc.UnaryServerInfo{
 		FullMethod: "/package.TestService.GetUser",
@@ -43,19 +48,73 @@ func TestUnaryOpenTracingInterceptor_Enabled(t *testing.T) {
 }
 
 func TestUnaryTracingInterceptor(t *testing.T) {
-	var run int32
-	var wg sync.WaitGroup
-	wg.Add(1)
-	_, err := UnaryTracingInterceptor(context.Background(), nil, &grpc.UnaryServerInfo{
-		FullMethod: "/",
-	}, func(ctx context.Context, req any) (any, error) {
-		defer wg.Done()
-		atomic.AddInt32(&run, 1)
-		return nil, nil
+	t.Run("normal", func(t *testing.T) {
+		var run int32
+		me := tracetest.NewInMemoryExporter(t)
+		_, err := UnaryTracingInterceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+			FullMethod: "/proto.Hello/Echo",
+		}, func(ctx context.Context, req any) (any, error) {
+			atomic.AddInt32(&run, 1)
+			return nil, nil
+		})
+		assert.Nil(t, err)
+		assert.Equal(t, int32(1), atomic.LoadInt32(&run))
+
+		assert.Equal(t, 1, len(me.GetSpans()))
+		span := me.GetSpans()[0].Snapshot()
+		assert.Equal(t, 2, len(span.Events()))
+		assert.ElementsMatch(t, []attribute.KeyValue{
+			ztrace.RPCSystemGRPC,
+			semconv.RPCServiceKey.String("proto.Hello"),
+			semconv.RPCMethodKey.String("Echo"),
+			ztrace.StatusCodeAttr(codes.OK),
+		}, span.Attributes())
+	})
+
+	t.Run("grpc error status", func(t *testing.T) {
+		me := tracetest.NewInMemoryExporter(t)
+		_, err := UnaryTracingInterceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+			FullMethod: "/proto.Hello/Echo",
+		}, func(ctx context.Context, req any) (any, error) {
+			return nil, status.Errorf(codes.Unknown, "test")
+		})
+		assert.Error(t, err)
+		assert.Equal(t, 1, len(me.GetSpans()))
+		span := me.GetSpans()[0].Snapshot()
+		assert.Equal(t, trace.Status{
+			Code:        tcodes.Error,
+			Description: "test",
+		}, span.Status())
+		assert.Equal(t, 2, len(span.Events()))
+		assert.ElementsMatch(t, []attribute.KeyValue{
+			ztrace.RPCSystemGRPC,
+			semconv.RPCServiceKey.String("proto.Hello"),
+			semconv.RPCMethodKey.String("Echo"),
+			ztrace.StatusCodeAttr(codes.Unknown),
+		}, span.Attributes())
+	})
+
+	t.Run("non grpc status error", func(t *testing.T) {
+		me := tracetest.NewInMemoryExporter(t)
+		_, err := UnaryTracingInterceptor(context.Background(), nil, &grpc.UnaryServerInfo{
+			FullMethod: "/proto.Hello/Echo",
+		}, func(ctx context.Context, req any) (any, error) {
+			return nil, errors.New("test")
+		})
+		assert.Error(t, err)
+		assert.Equal(t, 1, len(me.GetSpans()))
+		span := me.GetSpans()[0].Snapshot()
+		assert.Equal(t, trace.Status{
+			Code:        tcodes.Error,
+			Description: "test",
+		}, span.Status())
+		assert.Equal(t, 1, len(span.Events()))
+		assert.ElementsMatch(t, []attribute.KeyValue{
+			ztrace.RPCSystemGRPC,
+			semconv.RPCServiceKey.String("proto.Hello"),
+			semconv.RPCMethodKey.String("Echo"),
+		}, span.Attributes())
 	})
-	wg.Wait()
-	assert.Nil(t, err)
-	assert.Equal(t, int32(1), atomic.LoadInt32(&run))
 }
 
 func TestUnaryTracingInterceptor_WithError(t *testing.T) {