added select statement for case

grpc · Dec 4, 2024 · dea8408 · dea8408
1 parent 0cf3584
commit dea8408
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/internal/transport/handler_server_test.go b/internal/transport/handler_server_test.go
@@ -393,7 +393,6 @@ func (s) TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
 			t.Errorf("ctx.Err = %v; want %v", err, context.DeadlineExceeded)
 			return
 		}
-		s.WriteStatus(status.New(codes.DeadlineExceeded, "too slow"))
 		// rst flag setting to verify the noop function: signalDeadlineExceeded
 		ch := make(chan struct{}, 1)
 		origSignalDeadlineExceeded := signalDeadlineExceeded
@@ -403,9 +402,14 @@ func (s) TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
 		defer func() {
 			signalDeadlineExceeded = origSignalDeadlineExceeded
 		}()
-
+		s.WriteStatus(status.New(codes.DeadlineExceeded, "too slow"))
+		select {
+		case <-s.ctx.Done():
+		case <-time.After(5 * time.Second):
+			t.Errorf("timeout waiting for ctx.Done")
+			return
+		}
 	}
-
 	ht.HandleStreams(
 		context.Background(), func(s *ServerStream) { go runStream(s) },
 	)