@@ -358,6 +358,7 @@ public async Task Initialize_WhenNotSupportedCodec_ThrowWriterExceptionOnWriteAs
358
358
public async Task WriteAsync_WhenTransportExceptionOnWriteInWriterSession_ShouldReconnectAndReturnWriteResult ( )
359
359
{
360
360
var moveTcs = new TaskCompletionSource < bool > ( ) ;
361
+ var moveTcsRetry = new TaskCompletionSource < bool > ( ) ;
361
362
362
363
_mockStream . SetupSequence ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
363
364
. Returns ( Task . CompletedTask )
@@ -367,12 +368,17 @@ public async Task WriteAsync_WhenTransportExceptionOnWriteInWriterSession_Should
367
368
return new Driver . TransportException ( new RpcException ( Grpc . Core . Status . DefaultCancelled ) ) ;
368
369
} )
369
370
. Returns ( Task . CompletedTask )
370
- . Returns ( Task . CompletedTask ) ;
371
+ . Returns ( ( ) =>
372
+ {
373
+ moveTcsRetry . SetResult ( true ) ;
374
+
375
+ return Task . CompletedTask ;
376
+ } ) ;
371
377
_mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
372
378
. ReturnsAsync ( true )
373
379
. Returns ( new ValueTask < bool > ( moveTcs . Task ) )
374
380
. ReturnsAsync ( true )
375
- . ReturnsAsync ( true )
381
+ . Returns ( new ValueTask < bool > ( moveTcsRetry . Task ) )
376
382
. Returns ( _lastMoveNext ) ;
377
383
_mockStream . SetupSequence ( stream => stream . Current )
378
384
. Returns ( new StreamWriteMessage . Types . FromServer
@@ -569,7 +575,6 @@ public async Task WriteAsync_WhenStreamIsClosingOnProcessingWriteAck_ShouldRecon
569
575
public async Task WriteAsync_WhenCancellationTokenIsClosed_ThrowCancellationException ( )
570
576
{
571
577
var cancellationTokenSource = new CancellationTokenSource ( ) ;
572
- var nextCompleted = new TaskCompletionSource < bool > ( ) ;
573
578
_mockStream . Setup ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
574
579
. Returns ( Task . CompletedTask ) ;
575
580
_mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
@@ -582,10 +587,8 @@ public async Task WriteAsync_WhenCancellationTokenIsClosed_ThrowCancellationExce
582
587
583
588
var task = writer . WriteAsync ( 123L , cancellationTokenSource . Token ) ;
584
589
cancellationTokenSource . Cancel ( ) ;
585
- nextCompleted . SetResult ( true ) ;
586
590
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 ) ;
589
592
}
590
593
591
594
[ Fact ]
@@ -638,6 +641,7 @@ public async Task WriteAsync_WhenInFlightBufferSendInInitialize_ReturnCompletedT
638
641
var writeTcs2 = new TaskCompletionSource ( ) ;
639
642
var writeTcs3 = new TaskCompletionSource ( ) ;
640
643
var moveTcs = new TaskCompletionSource < bool > ( ) ;
644
+ var moveTcsRetry = new TaskCompletionSource < bool > ( ) ;
641
645
642
646
_mockStream . SetupSequence ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
643
647
. Returns ( Task . CompletedTask )
@@ -657,13 +661,17 @@ public async Task WriteAsync_WhenInFlightBufferSendInInitialize_ReturnCompletedT
657
661
return Task . CompletedTask ;
658
662
} )
659
663
. Returns ( Task . CompletedTask )
660
- . Returns ( Task . CompletedTask ) ;
664
+ . Returns ( ( ) =>
665
+ {
666
+ moveTcsRetry . SetResult ( true ) ;
667
+ return Task . CompletedTask ;
668
+ } ) ;
661
669
662
670
_mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
663
671
. ReturnsAsync ( true )
664
672
. Returns ( new ValueTask < bool > ( moveTcs . Task ) )
665
673
. ReturnsAsync ( true )
666
- . ReturnsAsync ( true )
674
+ . Returns ( new ValueTask < bool > ( moveTcsRetry . Task ) )
667
675
. Returns ( _lastMoveNext ) ;
668
676
_mockStream . SetupSequence ( stream => stream . Current )
669
677
. Returns ( new StreamWriteMessage . Types . FromServer
@@ -713,8 +721,7 @@ public async Task WriteAsync_WhenInFlightBufferSendInInitialize_ReturnCompletedT
713
721
714
722
moveTcs . SetResult ( false ) ; // Fail write ack stream => start reconnect
715
723
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 ) ;
718
725
Assert . Equal ( PersistenceStatus . AlreadyWritten , ( await runTask1 ) . Status ) ;
719
726
Assert . Equal ( PersistenceStatus . Written , ( await runTask2 ) . Status ) ;
720
727
@@ -863,6 +870,7 @@ public async Task DisposeAsync_WhenInFlightMessages_WaitingInFlightMessages()
863
870
{
864
871
var tcsDetectedWrite = new TaskCompletionSource ( ) ;
865
872
var writeTcs1 = new TaskCompletionSource < bool > ( ) ;
873
+ var moveTcsRetry = new TaskCompletionSource < bool > ( ) ;
866
874
867
875
_mockStream . SetupSequence ( stream => stream . Write ( It . IsAny < FromClient > ( ) ) )
868
876
. Returns ( Task . CompletedTask )
@@ -872,12 +880,16 @@ public async Task DisposeAsync_WhenInFlightMessages_WaitingInFlightMessages()
872
880
return Task . CompletedTask ;
873
881
} )
874
882
. Returns ( Task . CompletedTask )
875
- . Returns ( Task . CompletedTask ) ;
883
+ . Returns ( ( ) =>
884
+ {
885
+ moveTcsRetry . SetResult ( true ) ;
886
+ return Task . CompletedTask ;
887
+ } ) ;
876
888
_mockStream . SetupSequence ( stream => stream . MoveNextAsync ( ) )
877
889
. ReturnsAsync ( true )
878
890
. Returns ( new ValueTask < bool > ( writeTcs1 . Task ) )
879
891
. ReturnsAsync ( true )
880
- . ReturnsAsync ( true )
892
+ . Returns ( new ValueTask < bool > ( moveTcsRetry . Task ) )
881
893
. Returns ( _lastMoveNext ) ;
882
894
883
895
_mockStream . SetupSequence ( stream => stream . Current )
0 commit comments