1818import static com .google .common .truth .Truth .assertThat ;
1919
2020import com .google .api .core .SettableApiFuture ;
21+ import com .google .api .gax .batching .BatchEntry ;
2122import com .google .api .gax .batching .BatchingRequestBuilder ;
2223import com .google .api .gax .grpc .GrpcStatusCode ;
2324import com .google .api .gax .rpc .DeadlineExceededException ;
@@ -77,29 +78,43 @@ public void requestBuilderTest() {
7778
7879 @ Test
7980 public void splitResponseTest () {
80- List <SettableApiFuture <Void >> batchResponse =
81- ImmutableList .of (SettableApiFuture .<Void >create (), SettableApiFuture .<Void >create ());
82- assertThat (batchResponse .get (0 ).isDone ()).isFalse ();
83- assertThat (batchResponse .get (1 ).isDone ()).isFalse ();
81+ BatchEntry <RowMutationEntry , Void > batchEntry1 =
82+ BatchEntry .create (
83+ RowMutationEntry .create ("key1" ).deleteRow (), SettableApiFuture .<Void >create ());
84+ BatchEntry <RowMutationEntry , Void > batchEntry2 =
85+ BatchEntry .create (
86+ RowMutationEntry .create ("key2" ).deleteRow (), SettableApiFuture .<Void >create ());
87+
88+ List <BatchEntry <RowMutationEntry , Void >> batchResponse =
89+ ImmutableList .of (batchEntry1 , batchEntry2 );
90+ assertThat (batchResponse .get (0 ).getResultFuture ().isDone ()).isFalse ();
91+ assertThat (batchResponse .get (1 ).getResultFuture ().isDone ()).isFalse ();
8492
8593 MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor ();
8694 underTest .splitResponse (null , batchResponse );
87- assertThat (batchResponse .get (0 ).isDone ()).isTrue ();
88- assertThat (batchResponse .get (1 ).isDone ()).isTrue ();
95+ assertThat (batchResponse .get (0 ).getResultFuture (). isDone ()).isTrue ();
96+ assertThat (batchResponse .get (1 ).getResultFuture (). isDone ()).isTrue ();
8997 }
9098
9199 @ Test
92100 public void splitExceptionTest () {
101+ BatchEntry <RowMutationEntry , Void > batchEntry1 =
102+ BatchEntry .create (
103+ RowMutationEntry .create ("key1" ).deleteRow (), SettableApiFuture .<Void >create ());
104+ BatchEntry <RowMutationEntry , Void > batchEntry2 =
105+ BatchEntry .create (
106+ RowMutationEntry .create ("key2" ).deleteRow (), SettableApiFuture .<Void >create ());
107+
93108 MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor ();
94109 final RuntimeException expectedEx = new RuntimeException ("Caused while batching" );
95- List <SettableApiFuture < Void >> batchResponses =
96- ImmutableList .of (SettableApiFuture .< Void > create (), SettableApiFuture .< Void > create () );
110+ List <BatchEntry < RowMutationEntry , Void >> batchResponses =
111+ ImmutableList .of (batchEntry1 , batchEntry2 );
97112
98113 underTest .splitException (expectedEx , batchResponses );
99114
100- for (SettableApiFuture fu : batchResponses ) {
115+ for (BatchEntry < RowMutationEntry , Void > response : batchResponses ) {
101116 try {
102- fu .get ();
117+ response . getResultFuture () .get ();
103118 } catch (ExecutionException | InterruptedException ex ) {
104119 assertThat (ex ).hasCauseThat ().isSameInstanceAs (expectedEx );
105120 }
@@ -110,9 +125,15 @@ public void splitExceptionTest() {
110125 public void splitExceptionWithFailedMutationsTest () {
111126 MutateRowsBatchingDescriptor underTest = new MutateRowsBatchingDescriptor ();
112127 Throwable actualThrowable = null ;
113- SettableApiFuture <Void > resultFuture1 = SettableApiFuture .create ();
114- SettableApiFuture <Void > resultFuture2 = SettableApiFuture .create ();
115- SettableApiFuture <Void > resultFuture3 = SettableApiFuture .create ();
128+ BatchEntry <RowMutationEntry , Void > batchEntry1 =
129+ BatchEntry .create (
130+ RowMutationEntry .create ("key1" ).deleteRow (), SettableApiFuture .<Void >create ());
131+ BatchEntry <RowMutationEntry , Void > batchEntry2 =
132+ BatchEntry .create (
133+ RowMutationEntry .create ("key2" ).deleteRow (), SettableApiFuture .<Void >create ());
134+ BatchEntry <RowMutationEntry , Void > batchEntry3 =
135+ BatchEntry .create (
136+ RowMutationEntry .create ("key3" ).deleteRow (), SettableApiFuture .<Void >create ());
116137
117138 // Threw an exception at 1st and 3rd entry
118139 MutateRowsException serverError =
@@ -128,11 +149,10 @@ public void splitExceptionWithFailedMutationsTest() {
128149 new DeadlineExceededException (
129150 null , GrpcStatusCode .of (Status .Code .DEADLINE_EXCEEDED ), true ))),
130151 true );
131- underTest .splitException (
132- serverError , ImmutableList .of (resultFuture1 , resultFuture2 , resultFuture3 ));
152+ underTest .splitException (serverError , ImmutableList .of (batchEntry1 , batchEntry2 , batchEntry3 ));
133153
134154 try {
135- resultFuture1 .get ();
155+ batchEntry1 . getResultFuture () .get ();
136156 } catch (ExecutionException | InterruptedException e ) {
137157 actualThrowable = e ;
138158 }
@@ -143,15 +163,15 @@ public void splitExceptionWithFailedMutationsTest() {
143163 // As there is no exception for 2nd entry so it should not throw any exception
144164 actualThrowable = null ;
145165 try {
146- resultFuture2 .get ();
166+ batchEntry2 . getResultFuture () .get ();
147167 } catch (ExecutionException | InterruptedException e ) {
148168 actualThrowable = e ;
149169 }
150170 assertThat (actualThrowable ).isNull ();
151171
152172 actualThrowable = null ;
153173 try {
154- resultFuture3 .get ();
174+ batchEntry3 . getResultFuture () .get ();
155175 } catch (ExecutionException | InterruptedException e ) {
156176 actualThrowable = e ;
157177 }
0 commit comments