test: fix streaming unit tests (#1671)

cloudwego · Jan 14, 2025 · 99c5216 · 99c5216
1 parent c2b6b76
commit 99c5216
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/pkg/streamx/streamx_user_test.go b/pkg/streamx/streamx_user_test.go
@@ -803,13 +803,12 @@ func TestStreamingTracing(t *testing.T) {
 	test.Assert(t, err == nil, err)
 	defer svr.Stop()
 	tracer := &mockTracer{}
-	var finished sync.WaitGroup
+	finishedChan := make(chan struct{}, 1)
 	tracer.start = func(ctx context.Context) context.Context {
-		finished.Add(1)
 		return ctx
 	}
 	tracer.finish = func(ctx context.Context) {
-		finished.Done()
+		finishedChan <- struct{}{}
 	}
 	cli, err := NewClient(
 		"kitex.echo.service",
@@ -830,7 +829,7 @@ func TestStreamingTracing(t *testing.T) {
 		test.Assert(t, ri.Stats().Error() == nil, ri.Stats().Error())
 		_, err = stream.Recv(ctx)
 		test.Assert(t, err != nil, err)
-		finished.Wait()
+		<-finishedChan
 		// stream has been closed, Tracer.Finish should be called
 		// err set in rpcinfo.Stats() should be the same as the err returned to users
 		test.Assert(t, ri.Stats().Error() == err, ri.Stats().Error())
@@ -848,7 +847,7 @@ func TestStreamingTracing(t *testing.T) {
 		test.Assert(t, ri.Stats().Error() == nil, ri.Stats().Error())
 		_, err = stream.CloseAndRecv(ctx)
 		test.Assert(t, err != nil, err)
-		finished.Wait()
+		<-finishedChan
 		// stream has been closed, Tracer.Finish should be called
 		// err set in rpcinfo.Stats() should be the same as the err returned to users
 		test.Assert(t, ri.Stats().Error() == err, ri.Stats().Error())
@@ -868,7 +867,7 @@ func TestStreamingTracing(t *testing.T) {
 		test.Assert(t, err == nil, err)
 		_, err = stream.Recv(ctx)
 		test.Assert(t, err != nil, err)
-		finished.Wait()
+		<-finishedChan
 		// stream has been closed, Tracer.Finish should be called
 		// err set in rpcinfo.Stats() should be the same as the err returned to users
 		test.Assert(t, ri.Stats().Error() == err, ri.Stats().Error())