2323import static org .junit .Assert .assertThrows ;
2424
2525import com .google .api .gax .retrying .RetrySettings ;
26+ import com .google .api .gax .rpc .DeadlineExceededException ;
2627import com .google .api .gax .rpc .UnavailableException ;
2728import com .google .bigtable .v2 .BigtableGrpc ;
2829import com .google .bigtable .v2 .ExecuteQueryRequest ;
3637import com .google .cloud .bigtable .data .v2 .models .sql .Statement ;
3738import com .google .cloud .bigtable .data .v2 .stub .EnhancedBigtableStub ;
3839import com .google .cloud .bigtable .gaxx .testing .FakeStreamingApi .ServerStreamingStashCallable ;
39- import com .google .common .collect .Range ;
4040import io .grpc .Context ;
4141import io .grpc .Deadline ;
4242import io .grpc .Server ;
@@ -141,12 +141,13 @@ public void testExecuteQueryRequestsSetDefaultDeadline() {
141141 Iterator <SqlRow > iterator = stream .rows ().iterator ();
142142 // We don't care about this but are reusing the fake service that tests retries
143143 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 );
146147 }
147148
148149 @ Test
149- public void testExecuteQueryRequestsRespectOverridenDeadline () throws IOException {
150+ public void testExecuteQueryRequestsRespectDeadline () throws IOException {
150151 BigtableDataSettings .Builder overrideSettings =
151152 BigtableDataSettings .newBuilderForEmulator (server .getPort ())
152153 .setProjectId ("fake-project" )
@@ -156,18 +157,16 @@ public void testExecuteQueryRequestsRespectOverridenDeadline() throws IOExceptio
156157 .executeQuerySettings ()
157158 .setRetrySettings (
158159 RetrySettings .newBuilder ()
159- .setInitialRpcTimeout (Duration .ofMinutes ( 5 ))
160- .setMaxRpcTimeout (Duration .ofMinutes ( 5 ))
160+ .setInitialRpcTimeout (Duration .ofMillis ( 10 ))
161+ .setMaxRpcTimeout (Duration .ofMillis ( 10 ))
161162 .build ());
162163 EnhancedBigtableStub overrideDeadline =
163164 EnhancedBigtableStub .create (overrideSettings .build ().getStubSettings ());
164165 SqlServerStream streamOverride =
165166 overrideDeadline .executeQueryCallable ().call (Statement .of ("SELECT * FROM table" ));
166167 Iterator <SqlRow > overrideIterator = streamOverride .rows ().iterator ();
167168 // 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 ();
171170 }
172171
173172 private static class FakeService extends BigtableGrpc .BigtableImplBase {
@@ -181,6 +180,15 @@ public void executeQuery(
181180 Deadline deadline = Context .current ().getDeadline ();
182181 if (deadline != null ) {
183182 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 );
184192 }
185193 attempts ++;
186194 responseObserver .onNext (metadata (columnMetadata ("test" , stringType ())));
0 commit comments