Skip to content

Commit da00277

Browse files
rahul2393rahul2393
committed
fix tests
1 parent 1b00d27 commit da00277

2 files changed

Lines changed: 67 additions & 41 deletions

File tree

‎java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/LocationAwareSharedBackendReplicaHarnessTest.java‎

Lines changed: 60 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -105,12 +105,13 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
105105
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
106106

107107
seedLocationMetadata(client);
108-
waitForReplicaRoutedRead(client, harness, 0);
108+
int firstReplicaIndex = waitForReplicaRoutedRead(client, harness);
109+
int secondReplicaIndex = 1 - firstReplicaIndex;
109110
harness.clearRequests();
110111

111112
harness
112113
.replicas
113-
.get(0)
114+
.get(firstReplicaIndex)
114115
.putMethodErrors(
115116
SharedBackendReplicaHarness.METHOD_STREAMING_READ,
116117
resourceExhausted("busy-routed-replica"));
@@ -130,14 +131,14 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
130131
1,
131132
harness
132133
.replicas
133-
.get(0)
134+
.get(firstReplicaIndex)
134135
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
135136
.size());
136137
assertEquals(
137138
1,
138139
harness
139140
.replicas
140-
.get(1)
141+
.get(secondReplicaIndex)
141142
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
142143
.size());
143144
assertEquals(
@@ -150,19 +151,19 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
150151
(ReadRequest)
151152
harness
152153
.replicas
153-
.get(0)
154+
.get(firstReplicaIndex)
154155
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
155156
.get(0);
156157
assertTrue(replicaARequest.getResumeToken().isEmpty());
157158
assertRetriedOnSameLogicalRequest(
158159
harness
159160
.replicas
160-
.get(0)
161+
.get(firstReplicaIndex)
161162
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
162163
.get(0),
163164
harness
164165
.replicas
165-
.get(1)
166+
.get(secondReplicaIndex)
166167
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
167168
.get(0));
168169
}
@@ -176,12 +177,13 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
176177
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
177178

178179
seedLocationMetadata(client);
179-
waitForReplicaRoutedRead(client, harness, 0);
180+
int firstReplicaIndex = waitForReplicaRoutedRead(client, harness);
181+
int secondReplicaIndex = 1 - firstReplicaIndex;
180182
harness.clearRequests();
181183

182184
harness
183185
.replicas
184-
.get(0)
186+
.get(firstReplicaIndex)
185187
.putMethodErrors(
186188
SharedBackendReplicaHarness.METHOD_STREAMING_READ,
187189
resourceExhaustedWithRetryInfo("busy-routed-replica"));
@@ -212,14 +214,14 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
212214
1,
213215
harness
214216
.replicas
215-
.get(0)
217+
.get(firstReplicaIndex)
216218
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
217219
.size());
218220
assertEquals(
219221
2,
220222
harness
221223
.replicas
222-
.get(1)
224+
.get(secondReplicaIndex)
223225
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
224226
.size());
225227
assertEquals(
@@ -229,16 +231,22 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
229231
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
230232
.size());
231233
List<AbstractMessage> replicaBRequests =
232-
harness.replicas.get(1).getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
234+
harness
235+
.replicas
236+
.get(secondReplicaIndex)
237+
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
233238
for (AbstractMessage request : replicaBRequests) {
234239
assertTrue(((ReadRequest) request).getResumeToken().isEmpty());
235240
}
236241
List<String> replicaBRequestIds =
237-
harness.replicas.get(1).getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
242+
harness
243+
.replicas
244+
.get(secondReplicaIndex)
245+
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
238246
assertRetriedOnSameLogicalRequest(
239247
harness
240248
.replicas
241-
.get(0)
249+
.get(firstReplicaIndex)
242250
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
243251
.get(0),
244252
replicaBRequestIds.get(0));
@@ -256,12 +264,13 @@ public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exceptio
256264
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
257265

258266
seedLocationMetadata(client);
259-
waitForReplicaRoutedRead(client, harness, 0);
267+
int firstReplicaIndex = waitForReplicaRoutedRead(client, harness);
268+
int secondReplicaIndex = 1 - firstReplicaIndex;
260269
harness.clearRequests();
261270

262271
harness
263272
.replicas
264-
.get(0)
273+
.get(firstReplicaIndex)
265274
.putMethodErrors(
266275
SharedBackendReplicaHarness.METHOD_STREAMING_READ, unavailable("isolated-replica"));
267276

@@ -280,14 +289,14 @@ public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exceptio
280289
1,
281290
harness
282291
.replicas
283-
.get(0)
292+
.get(firstReplicaIndex)
284293
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
285294
.size());
286295
assertEquals(
287296
1,
288297
harness
289298
.replicas
290-
.get(1)
299+
.get(secondReplicaIndex)
291300
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
292301
.size());
293302
assertEquals(
@@ -300,19 +309,19 @@ public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exceptio
300309
(ReadRequest)
301310
harness
302311
.replicas
303-
.get(0)
312+
.get(firstReplicaIndex)
304313
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
305314
.get(0);
306315
assertTrue(replicaARequest.getResumeToken().isEmpty());
307316
assertRetriedOnSameLogicalRequest(
308317
harness
309318
.replicas
310-
.get(0)
319+
.get(firstReplicaIndex)
311320
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
312321
.get(0),
313322
harness
314323
.replicas
315-
.get(1)
324+
.get(secondReplicaIndex)
316325
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
317326
.get(0));
318327
}
@@ -327,12 +336,13 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
327336
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
328337

329338
seedLocationMetadata(client);
330-
waitForReplicaRoutedRead(client, harness, 0);
339+
int firstReplicaIndex = waitForReplicaRoutedRead(client, harness);
340+
int secondReplicaIndex = 1 - firstReplicaIndex;
331341
harness.clearRequests();
332342

333343
harness
334344
.replicas
335-
.get(0)
345+
.get(firstReplicaIndex)
336346
.putMethodErrors(
337347
SharedBackendReplicaHarness.METHOD_STREAMING_READ, unavailable("isolated-replica"));
338348

@@ -362,14 +372,14 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
362372
1,
363373
harness
364374
.replicas
365-
.get(0)
375+
.get(firstReplicaIndex)
366376
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
367377
.size());
368378
assertEquals(
369379
2,
370380
harness
371381
.replicas
372-
.get(1)
382+
.get(secondReplicaIndex)
373383
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
374384
.size());
375385
assertEquals(
@@ -379,16 +389,22 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
379389
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
380390
.size());
381391
List<AbstractMessage> replicaBRequests =
382-
harness.replicas.get(1).getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
392+
harness
393+
.replicas
394+
.get(secondReplicaIndex)
395+
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
383396
for (AbstractMessage request : replicaBRequests) {
384397
assertTrue(((ReadRequest) request).getResumeToken().isEmpty());
385398
}
386399
List<String> replicaBRequestIds =
387-
harness.replicas.get(1).getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
400+
harness
401+
.replicas
402+
.get(secondReplicaIndex)
403+
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ);
388404
assertRetriedOnSameLogicalRequest(
389405
harness
390406
.replicas
391-
.get(0)
407+
.get(firstReplicaIndex)
392408
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
393409
.get(0),
394410
replicaBRequestIds.get(0));
@@ -407,7 +423,8 @@ public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTra
407423
DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
408424

409425
seedLocationMetadata(client);
410-
waitForReplicaRoutedRead(client, harness, 0);
426+
int firstReplicaIndex = waitForReplicaRoutedRead(client, harness);
427+
int secondReplicaIndex = 1 - firstReplicaIndex;
411428
harness.clearRequests();
412429

413430
harness.backend.setStreamingReadExecutionTime(
@@ -432,14 +449,14 @@ public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTra
432449
1,
433450
harness
434451
.replicas
435-
.get(0)
452+
.get(firstReplicaIndex)
436453
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
437454
.size());
438455
assertEquals(
439456
1,
440457
harness
441458
.replicas
442-
.get(1)
459+
.get(secondReplicaIndex)
443460
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
444461
.size());
445462
assertEquals(
@@ -453,27 +470,27 @@ public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTra
453470
(ReadRequest)
454471
harness
455472
.replicas
456-
.get(0)
473+
.get(firstReplicaIndex)
457474
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
458475
.get(0);
459476
ReadRequest replicaBRequest =
460477
(ReadRequest)
461478
harness
462479
.replicas
463-
.get(1)
480+
.get(secondReplicaIndex)
464481
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
465482
.get(0);
466483
assertTrue(replicaARequest.getResumeToken().isEmpty());
467484
assertEquals(RESUME_TOKEN_AFTER_FIRST_ROW, replicaBRequest.getResumeToken());
468485
assertRetriedOnSameLogicalRequest(
469486
harness
470487
.replicas
471-
.get(0)
488+
.get(firstReplicaIndex)
472489
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
473490
.get(0),
474491
harness
475492
.replicas
476-
.get(1)
493+
.get(secondReplicaIndex)
477494
.getRequestIds(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
478495
.get(0));
479496
}
@@ -514,9 +531,8 @@ private static void seedLocationMetadata(DatabaseClient client) {
514531
}
515532
}
516533

517-
private static void waitForReplicaRoutedRead(
518-
DatabaseClient client, SharedBackendReplicaHarness harness, int replicaIndex)
519-
throws InterruptedException {
534+
private static int waitForReplicaRoutedRead(
535+
DatabaseClient client, SharedBackendReplicaHarness harness) throws InterruptedException {
520536
long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(10);
521537
while (System.nanoTime() < deadlineNanos) {
522538
try (ResultSet resultSet =
@@ -527,13 +543,16 @@ private static void waitForReplicaRoutedRead(
527543
KeySet.singleKey(Key.of("b")),
528544
Arrays.asList("k"),
529545
Options.directedRead(DIRECTED_READ_OPTIONS))) {
530-
if (resultSet.next()
531-
&& !harness
546+
if (resultSet.next()) {
547+
for (int replicaIndex = 0; replicaIndex < harness.replicas.size(); replicaIndex++) {
548+
if (!harness
532549
.replicas
533550
.get(replicaIndex)
534551
.getRequests(SharedBackendReplicaHarness.METHOD_STREAMING_READ)
535552
.isEmpty()) {
536-
return;
553+
return replicaIndex;
554+
}
555+
}
537556
}
538557
}
539558
Thread.sleep(50L);

‎java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/spi/v1/KeyAwareChannelTest.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@
6666
import java.util.Map;
6767
import java.util.concurrent.TimeUnit;
6868
import javax.annotation.Nullable;
69+
import org.junit.After;
6970
import org.junit.Test;
7071
import org.junit.runner.RunWith;
7172
import org.junit.runners.JUnit4;
@@ -76,6 +77,12 @@ public class KeyAwareChannelTest {
7677
private static final String SESSION =
7778
"projects/p/instances/i/databases/d/sessions/test-session-id";
7879

80+
@After
81+
public void clearSharedRoutingState() {
82+
EndpointLatencyRegistry.clear();
83+
RequestIdTargetTracker.clear();
84+
}
85+
7986
@Test
8087
public void cancelBeforeStartPreservesTrailersAndSkipsDelegateCreation() throws Exception {
8188
TestHarness harness = createHarness();

0 commit comments

Comments
 (0)