From dea8408cd265d1302572465211293f8f727ab9ef Mon Sep 17 00:00:00 2001 From: misvivek Date: Wed, 4 Dec 2024 11:45:11 +0000 Subject: [PATCH] added select statement for case --- internal/transport/handler_server_test.go | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/internal/transport/handler_server_test.go b/internal/transport/handler_server_test.go index c17aad611198..8ac3734fb200 100644 --- 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) }, )