23
23
import static org .junit .Assert .assertThrows ;
24
24
25
25
import com .google .api .gax .retrying .RetrySettings ;
26
+ import com .google .api .gax .rpc .DeadlineExceededException ;
26
27
import com .google .api .gax .rpc .UnavailableException ;
27
28
import com .google .bigtable .v2 .BigtableGrpc ;
28
29
import com .google .bigtable .v2 .ExecuteQueryRequest ;
36
37
import com .google .cloud .bigtable .data .v2 .models .sql .Statement ;
37
38
import com .google .cloud .bigtable .data .v2 .stub .EnhancedBigtableStub ;
38
39
import com .google .cloud .bigtable .gaxx .testing .FakeStreamingApi .ServerStreamingStashCallable ;
39
- import com .google .common .collect .Range ;
40
40
import io .grpc .Context ;
41
41
import io .grpc .Deadline ;
42
42
import io .grpc .Server ;
@@ -141,12 +141,13 @@ public void testExecuteQueryRequestsSetDefaultDeadline() {
141
141
Iterator <SqlRow > iterator = stream .rows ().iterator ();
142
142
// We don't care about this but are reusing the fake service that tests retries
143
143
assertThrows (UnavailableException .class , iterator ::next ).getCause ();
144
- // We have 30s default, we assume less than 1s has been burned when the fake service sets it
145
- assertThat (fakeService .deadlineMillisRemaining ).isIn (Range .closed (29000L , 30100L ));
144
+ // We have 30s default, we give it a wide range to avoid flakiness, this is mostly just checking
145
+ // that some default is set
146
+ assertThat (fakeService .deadlineMillisRemaining ).isLessThan (30001L );
146
147
}
147
148
148
149
@ Test
149
- public void testExecuteQueryRequestsRespectOverridenDeadline () throws IOException {
150
+ public void testExecuteQueryRequestsRespectDeadline () throws IOException {
150
151
BigtableDataSettings .Builder overrideSettings =
151
152
BigtableDataSettings .newBuilderForEmulator (server .getPort ())
152
153
.setProjectId ("fake-project" )
@@ -156,18 +157,16 @@ public void testExecuteQueryRequestsRespectOverridenDeadline() throws IOExceptio
156
157
.executeQuerySettings ()
157
158
.setRetrySettings (
158
159
RetrySettings .newBuilder ()
159
- .setInitialRpcTimeout (Duration .ofMinutes ( 5 ))
160
- .setMaxRpcTimeout (Duration .ofMinutes ( 5 ))
160
+ .setInitialRpcTimeout (Duration .ofMillis ( 10 ))
161
+ .setMaxRpcTimeout (Duration .ofMillis ( 10 ))
161
162
.build ());
162
163
EnhancedBigtableStub overrideDeadline =
163
164
EnhancedBigtableStub .create (overrideSettings .build ().getStubSettings ());
164
165
SqlServerStream streamOverride =
165
166
overrideDeadline .executeQueryCallable ().call (Statement .of ("SELECT * FROM table" ));
166
167
Iterator <SqlRow > overrideIterator = streamOverride .rows ().iterator ();
167
168
// We don't care about this but are reusing the fake service that tests retries
168
- assertThrows (UnavailableException .class , overrideIterator ::next ).getCause ();
169
- // We have 30s default, we assume less than 1s has been burned when the fake service sets it
170
- assertThat (fakeService .deadlineMillisRemaining ).isIn (Range .closed (299000L , 300100L ));
169
+ assertThrows (DeadlineExceededException .class , overrideIterator ::next ).getCause ();
171
170
}
172
171
173
172
private static class FakeService extends BigtableGrpc .BigtableImplBase {
@@ -181,6 +180,15 @@ public void executeQuery(
181
180
Deadline deadline = Context .current ().getDeadline ();
182
181
if (deadline != null ) {
183
182
deadlineMillisRemaining = deadline .timeRemaining (TimeUnit .MILLISECONDS );
183
+ } else {
184
+ // set to max long when deadline isn't set
185
+ deadlineMillisRemaining = Long .MAX_VALUE ;
186
+ }
187
+ // Sleep for 100ms to trigger deadline exceeded for tests with a shorter deadline
188
+ try {
189
+ Thread .sleep (100 );
190
+ } catch (InterruptedException e ) {
191
+ throw new RuntimeException (e );
184
192
}
185
193
attempts ++;
186
194
responseObserver .onNext (metadata (columnMetadata ("test" , stringType ())));
0 commit comments