3333import com .google .api .gax .rpc .ResponseObserver ;
3434import com .google .api .gax .rpc .StatusCode ;
3535import com .google .api .gax .rpc .StreamController ;
36+ import com .google .api .gax .rpc .UnaryCallable ;
3637import com .google .cloud .pubsub .v1 .stub .SubscriberStub ;
3738import com .google .common .collect .Lists ;
3839import com .google .protobuf .Any ;
40+ import com .google .protobuf .Empty ;
3941import com .google .pubsub .v1 .AcknowledgeRequest ;
4042import com .google .pubsub .v1 .ModifyAckDeadlineRequest ;
4143import com .google .pubsub .v1 .StreamingPullRequest ;
@@ -66,6 +68,9 @@ public class StreamingSubscriberConnectionTest {
6668 private FakeScheduledExecutorService executor ;
6769 private FakeClock clock ;
6870 private SubscriberStub mockSubscriberStub ;
71+ private BidiStreamingCallable <StreamingPullRequest , StreamingPullResponse > mockStreamingCallable ;
72+ private UnaryCallable <AcknowledgeRequest , Empty > mockAcknowledgeCallable ;
73+ private UnaryCallable <ModifyAckDeadlineRequest , Empty > mockModifyAckDeadlineCallable ;
6974
7075 private static final String MOCK_SUBSCRIPTION_NAME =
7176 "projects/MOCK-PROJECT/subscriptions/MOCK-SUBSCRIPTION" ;
@@ -99,14 +104,28 @@ public class StreamingSubscriberConnectionTest {
99104 private static final Duration MAX_ACK_EXTENSION_PERIOD = Duration .ofMinutes (60 );
100105
101106 @ Before
102- public void setUp () {
107+ public void setUp () throws Exception {
103108 systemExecutor = new FakeScheduledExecutorService ();
104109 executor = new FakeScheduledExecutorService ();
105110 clock = systemExecutor .getClock ();
106- mockSubscriberStub =
107- mock (
108- SubscriberStub .class ,
109- withSettings ().withoutAnnotations ().defaultAnswer (RETURNS_DEEP_STUBS ));
111+
112+ mockStreamingCallable = mock (BidiStreamingCallable .class , withSettings ().withoutAnnotations ());
113+ mockAcknowledgeCallable = mock (UnaryCallable .class , withSettings ().withoutAnnotations ());
114+ mockModifyAckDeadlineCallable = mock (UnaryCallable .class , withSettings ().withoutAnnotations ());
115+
116+ mockSubscriberStub = mock (SubscriberStub .class , withSettings ().withoutAnnotations ());
117+ when (mockSubscriberStub .streamingPullCallable ()).thenReturn (mockStreamingCallable );
118+ when (mockSubscriberStub .acknowledgeCallable ()).thenReturn (mockAcknowledgeCallable );
119+ when (mockSubscriberStub .modifyAckDeadlineCallable ()).thenReturn (mockModifyAckDeadlineCallable );
120+
121+ ClientStream <StreamingPullRequest > defaultMockClientStream =
122+ mock (ClientStream .class , withSettings ().withoutAnnotations ());
123+ when (mockStreamingCallable .splitCall (any (), any ())).thenReturn (defaultMockClientStream );
124+
125+ when (mockAcknowledgeCallable .futureCall (any ()))
126+ .thenAnswer (invocation -> SettableApiFuture .create ());
127+ when (mockModifyAckDeadlineCallable .futureCall (any ()))
128+ .thenAnswer (invocation -> SettableApiFuture .create ());
110129 }
111130
112131 @ After
@@ -391,12 +410,9 @@ public void testSendAckOperationsExactlyOnceEnabledMessageFuturesModacks() {
391410 systemExecutor .advanceTime (Duration .ofSeconds (200 ));
392411
393412 // Assert expected behavior
394- verify (mockSubscriberStub .modifyAckDeadlineCallable (), times (1 ))
395- .futureCall (modifyAckDeadlineRequestNack );
396- verify (mockSubscriberStub .modifyAckDeadlineCallable (), times (1 ))
397- .futureCall (modifyAckDeadlineRequestInitial );
398- verify (mockSubscriberStub .modifyAckDeadlineCallable (), times (1 ))
399- .futureCall (modifyAckDeadlineRequestRetry );
413+ verify (mockModifyAckDeadlineCallable , times (1 )).futureCall (modifyAckDeadlineRequestNack );
414+ verify (mockModifyAckDeadlineCallable , times (1 )).futureCall (modifyAckDeadlineRequestInitial );
415+ verify (mockModifyAckDeadlineCallable , times (1 )).futureCall (modifyAckDeadlineRequestRetry );
400416 verify (mockSubscriberStub , never ()).acknowledgeCallable ();
401417
402418 try {
@@ -511,9 +527,8 @@ public void testSendAckOperationsExactlyOnceEnabledMessageFuturesAcks() {
511527 systemExecutor .advanceTime (Duration .ofMillis (200 ));
512528
513529 // Assert expected behavior;
514- verify (mockSubscriberStub .acknowledgeCallable (), times (1 ))
515- .futureCall (acknowledgeRequestInitial );
516- verify (mockSubscriberStub .acknowledgeCallable (), times (1 ))
530+ verify (mockAcknowledgeCallable , times (1 )).futureCall (acknowledgeRequestInitial );
531+ verify (mockAcknowledgeCallable , times (1 ))
517532 .futureCall (
518533 argThat (new CustomArgumentMatchers .AcknowledgeRequestMatcher (acknowledgeRequestRetry )));
519534 verify (mockSubscriberStub , never ()).modifyAckDeadlineCallable ();
@@ -577,7 +592,7 @@ public void testSendAckOperationsExactlyOnceEnabledErrorWithEmptyMetadataMap() {
577592 systemExecutor .advanceTime (Duration .ofMillis (200 ));
578593
579594 // Assert expected behavior;
580- verify (mockSubscriberStub . acknowledgeCallable () , times (2 )).futureCall (acknowledgeRequest );
595+ verify (mockAcknowledgeCallable , times (2 )).futureCall (acknowledgeRequest );
581596 verify (mockSubscriberStub , never ()).modifyAckDeadlineCallable ();
582597
583598 try {
@@ -677,27 +692,22 @@ public void testMaxPerRequestChanges() {
677692 .setSubscription (MOCK_SUBSCRIPTION_NAME )
678693 .addAllAckIds (mockAckIdsInRequest )
679694 .build ();
680- verify (mockSubscriberStub .acknowledgeCallable (), times (1 ))
681- .futureCall (expectedAcknowledgeRequest );
695+ verify (mockAcknowledgeCallable , times (1 )).futureCall (expectedAcknowledgeRequest );
682696
683697 ModifyAckDeadlineRequest expectedModifyAckDeadlineRequest =
684698 ModifyAckDeadlineRequest .newBuilder ()
685699 .setSubscription (MOCK_SUBSCRIPTION_NAME )
686700 .addAllAckIds (mockAckIdsInRequest )
687701 .setAckDeadlineSeconds (MOCK_ACK_EXTENSION_DEFAULT_SECONDS )
688702 .build ();
689- verify (mockSubscriberStub .modifyAckDeadlineCallable (), times (1 ))
690- .futureCall (expectedModifyAckDeadlineRequest );
703+ verify (mockModifyAckDeadlineCallable , times (1 )).futureCall (expectedModifyAckDeadlineRequest );
691704 }
692705 }
693706
694707 @ Test
695708 public void testClientPinger_pingSent () {
696- BidiStreamingCallable <StreamingPullRequest , StreamingPullResponse > mockStreamingCallable =
697- mock (BidiStreamingCallable .class , withSettings ().withoutAnnotations ());
698709 ClientStream <StreamingPullRequest > mockClientStream =
699710 mock (ClientStream .class , withSettings ().withoutAnnotations ());
700- when (mockSubscriberStub .streamingPullCallable ()).thenReturn (mockStreamingCallable );
701711 when (mockStreamingCallable .splitCall (any (ResponseObserver .class ), any ()))
702712 .thenReturn (mockClientStream );
703713
@@ -737,11 +747,8 @@ public void testClientPinger_pingSent() {
737747
738748 @ Test
739749 public void testClientPinger_pingsNotSentWhenDisabled () {
740- BidiStreamingCallable <StreamingPullRequest , StreamingPullResponse > mockStreamingCallable =
741- mock (BidiStreamingCallable .class , withSettings ().withoutAnnotations ());
742750 ClientStream <StreamingPullRequest > mockClientStream =
743751 mock (ClientStream .class , withSettings ().withoutAnnotations ());
744- when (mockSubscriberStub .streamingPullCallable ()).thenReturn (mockStreamingCallable );
745752 when (mockStreamingCallable .splitCall (any (ResponseObserver .class ), any ()))
746753 .thenReturn (mockClientStream );
747754
@@ -763,13 +770,10 @@ public void testClientPinger_pingsNotSentWhenDisabled() {
763770
764771 @ Test
765772 public void testServerMonitor_timesOut () {
766- BidiStreamingCallable <StreamingPullRequest , StreamingPullResponse > mockStreamingCallable =
767- mock (BidiStreamingCallable .class , withSettings ().withoutAnnotations ());
768773 ClientStream <StreamingPullRequest > mockClientStream =
769774 mock (ClientStream .class , withSettings ().withoutAnnotations ());
770775 ArgumentCaptor <ResponseObserver <StreamingPullResponse >> observerCaptor =
771776 ArgumentCaptor .forClass (ResponseObserver .class );
772- when (mockSubscriberStub .streamingPullCallable ()).thenReturn (mockStreamingCallable );
773777 when (mockStreamingCallable .splitCall (observerCaptor .capture (), any ()))
774778 .thenReturn (mockClientStream );
775779
@@ -813,13 +817,10 @@ public void testServerMonitor_timesOut() {
813817
814818 @ Test
815819 public void testServerMonitor_doesNotTimeOutIfResponseReceived () {
816- BidiStreamingCallable <StreamingPullRequest , StreamingPullResponse > mockStreamingCallable =
817- mock (BidiStreamingCallable .class , withSettings ().withoutAnnotations ());
818820 ClientStream <StreamingPullRequest > mockClientStream =
819821 mock (ClientStream .class , withSettings ().withoutAnnotations ());
820822 ArgumentCaptor <ResponseObserver <StreamingPullResponse >> observerCaptor =
821823 ArgumentCaptor .forClass (ResponseObserver .class );
822- when (mockSubscriberStub .streamingPullCallable ()).thenReturn (mockStreamingCallable );
823824 when (mockStreamingCallable .splitCall (observerCaptor .capture (), any ()))
824825 .thenReturn (mockClientStream );
825826
0 commit comments