change done on checkHeaderandtrailer

internal/transport/handler_server_test.go

@@ -380,6 +380,15 @@ func (s) TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	// rst flag setting to verify the noop function: signalDeadlineExceeded
+	ch := make(chan struct{}, 1)
+	origSignalDeadlineExceeded := signalDeadlineExceeded
+	signalDeadlineExceeded = func() {
+		ch <- struct{}{}
+	}
+	defer func() {
+		signalDeadlineExceeded = origSignalDeadlineExceeded
+	}()
 	runStream := func(s *ServerStream) {
 		defer bodyw.Close()
 		select {
@@ -393,23 +402,9 @@ func (s) TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
 			t.Errorf("ctx.Err = %v; want %v", err, context.DeadlineExceeded)
 			return
 		}
-		// rst flag setting to verify the noop function: signalDeadlineExceeded
-		ch := make(chan struct{}, 1)
-		origSignalDeadlineExceeded := signalDeadlineExceeded
-		signalDeadlineExceeded = func() {
-			ch <- struct{}{}
-		}
-		defer func() {
-			signalDeadlineExceeded = origSignalDeadlineExceeded
-		}()
+
 		s.WriteStatus(status.New(codes.DeadlineExceeded, "too slow"))
-		select {
-		case <-ch: // Signal received, continue with the test
-		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) },
@@ -424,6 +419,13 @@ func (s) TestHandlerTransport_HandleStreams_Timeout(t *testing.T) {
 		"Grpc-Message": {encodeGrpcMessage("too slow")},
 	}
 	checkHeaderAndTrailer(t, rw, wantHeader, wantTrailer)
+	select {
+	case <-ch: // Signal received, continue with the test
+	case <-time.After(5 * time.Second):
+		t.Errorf("timeout waiting for ctx.Done")
+		return
+	}
+
 }
 
 // TestHandlerTransport_HandleStreams_MultiWriteStatus ensures that