@@ -358,6 +358,7 @@ public async Task Initialize_WhenNotSupportedCodec_ThrowWriterExceptionOnWriteAs
358358 public async Task WriteAsync_WhenTransportExceptionOnWriteInWriterSession_ShouldReconnectAndReturnWriteResult ( )
359359 {
360360 var moveTcs = new TaskCompletionSource < bool > ( ) ;
361+ var moveTcsRetry = new TaskCompletionSource < bool > ( ) ;
361362
362363 _mockStream . SetupSequence ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
363364 . Returns ( Task . CompletedTask )
@@ -367,12 +368,17 @@ public async Task WriteAsync_WhenTransportExceptionOnWriteInWriterSession_Should
367368 return new Driver . TransportException ( new RpcException ( Grpc . Core . Status . DefaultCancelled ) ) ;
368369 } )
369370 . Returns ( Task . CompletedTask )
370- . Returns ( Task . CompletedTask ) ;
371+ . Returns ( ( ) =>
372+ {
373+ moveTcsRetry . SetResult ( true ) ;
374+
375+ return Task . CompletedTask ;
376+ } ) ;
371377 _mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
372378 . ReturnsAsync ( true )
373379 . Returns ( new ValueTask < bool > ( moveTcs . Task ) )
374380 . ReturnsAsync ( true )
375- . ReturnsAsync ( true )
381+ . Returns ( new ValueTask < bool > ( moveTcsRetry . Task ) )
376382 . Returns ( _lastMoveNext ) ;
377383 _mockStream . SetupSequence ( stream => stream . Current )
378384 . Returns ( new StreamWriteMessage . Types . FromServer
@@ -569,7 +575,6 @@ public async Task WriteAsync_WhenStreamIsClosingOnProcessingWriteAck_ShouldRecon
569575 public async Task WriteAsync_WhenCancellationTokenIsClosed_ThrowCancellationException ( )
570576 {
571577 var cancellationTokenSource = new CancellationTokenSource ( ) ;
572- var nextCompleted = new TaskCompletionSource < bool > ( ) ;
573578 _mockStream . Setup ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
574579 . Returns ( Task . CompletedTask ) ;
575580 _mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
@@ -582,10 +587,8 @@ public async Task WriteAsync_WhenCancellationTokenIsClosed_ThrowCancellationExce
582587
583588 var task = writer . WriteAsync ( 123L , cancellationTokenSource . Token ) ;
584589 cancellationTokenSource . Cancel ( ) ;
585- nextCompleted . SetResult ( true ) ;
586590
587- Assert . Equal ( "The write operation was canceled before it could be completed" ,
588- ( await Assert . ThrowsAsync < WriterException > ( ( ) => task ) ) . Message ) ;
591+ await Assert . ThrowsAsync < TaskCanceledException > ( ( ) => task ) ;
589592 }
590593
591594 [ Fact ]
@@ -638,6 +641,7 @@ public async Task WriteAsync_WhenInFlightBufferSendInInitialize_ReturnCompletedT
638641 var writeTcs2 = new TaskCompletionSource ( ) ;
639642 var writeTcs3 = new TaskCompletionSource ( ) ;
640643 var moveTcs = new TaskCompletionSource < bool > ( ) ;
644+ var moveTcsRetry = new TaskCompletionSource < bool > ( ) ;
641645
642646 _mockStream . SetupSequence ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
643647 . Returns ( Task . CompletedTask )
@@ -657,13 +661,17 @@ public async Task WriteAsync_WhenInFlightBufferSendInInitialize_ReturnCompletedT
657661 return Task . CompletedTask ;
658662 } )
659663 . Returns ( Task . CompletedTask )
660- . Returns ( Task . CompletedTask ) ;
664+ . Returns ( ( ) =>
665+ {
666+ moveTcsRetry . SetResult ( true ) ;
667+ return Task . CompletedTask ;
668+ } ) ;
661669
662670 _mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
663671 . ReturnsAsync ( true )
664672 . Returns ( new ValueTask < bool > ( moveTcs . Task ) )
665673 . ReturnsAsync ( true )
666- . ReturnsAsync ( true )
674+ . Returns ( new ValueTask < bool > ( moveTcsRetry . Task ) )
667675 . Returns ( _lastMoveNext ) ;
668676 _mockStream . SetupSequence ( stream => stream . Current )
669677 . Returns ( new StreamWriteMessage . Types . FromServer
@@ -713,8 +721,7 @@ public async Task WriteAsync_WhenInFlightBufferSendInInitialize_ReturnCompletedT
713721
714722 moveTcs . SetResult ( false ) ; // Fail write ack stream => start reconnect
715723
716- Assert . Equal ( "The write operation was canceled before it could be completed" ,
717- ( await Assert . ThrowsAsync < WriterException > ( ( ) => runTaskWithCancel ) ) . Message ) ;
724+ await Assert . ThrowsAsync < TaskCanceledException > ( ( ) => runTaskWithCancel ) ;
718725 Assert . Equal ( PersistenceStatus . AlreadyWritten , ( await runTask1 ) . Status ) ;
719726 Assert . Equal ( PersistenceStatus . Written , ( await runTask2 ) . Status ) ;
720727
@@ -863,6 +870,7 @@ public async Task DisposeAsync_WhenInFlightMessages_WaitingInFlightMessages()
863870 {
864871 var tcsDetectedWrite = new TaskCompletionSource ( ) ;
865872 var writeTcs1 = new TaskCompletionSource < bool > ( ) ;
873+ var moveTcsRetry = new TaskCompletionSource < bool > ( ) ;
866874
867875 _mockStream . SetupSequence ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
868876 . Returns ( Task . CompletedTask )
@@ -872,12 +880,16 @@ public async Task DisposeAsync_WhenInFlightMessages_WaitingInFlightMessages()
872880 return Task . CompletedTask ;
873881 } )
874882 . Returns ( Task . CompletedTask )
875- . Returns ( Task . CompletedTask ) ;
883+ . Returns ( ( ) =>
884+ {
885+ moveTcsRetry . SetResult ( true ) ;
886+ return Task . CompletedTask ;
887+ } ) ;
876888 _mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
877889 . ReturnsAsync ( true )
878890 . Returns ( new ValueTask < bool > ( writeTcs1 . Task ) )
879891 . ReturnsAsync ( true )
880- . ReturnsAsync ( true )
892+ . Returns ( new ValueTask < bool > ( moveTcsRetry . Task ) )
881893 . Returns ( _lastMoveNext ) ;
882894
883895 _mockStream . SetupSequence ( stream => stream . Current )
0 commit comments