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

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) {