speed up tests

Author: rahul2393

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

@@ -55,6 +55,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.AfterClass;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -82,27 +83,40 @@ public class LocationAwareSharedBackendReplicaHarnessTest {
                           .build())
                   .build())
           .build();
+  private static SharedBackendReplicaHarness harness;
 
   @BeforeClass
-  public static void enableLocationAwareRouting() {
+  public static void enableLocationAwareRouting() throws Exception {
     SpannerOptions.useEnvironment(
         new SpannerOptions.SpannerEnvironment() {
           @Override
           public boolean isEnableLocationApi() {
             return true;
           }
         });
+    harness = SharedBackendReplicaHarness.create(2);
+  }
+
+  @Before
+  public void resetHarness() {
+    harness.reset();
   }
 
   @AfterClass
-  public static void restoreEnvironment() {
-    SpannerOptions.useDefaultEnvironment();
+  public static void restoreEnvironment() throws Exception {
+    try {
+      if (harness != null) {
+        harness.close();
+      }
+    } finally {
+      harness = null;
+      SpannerOptions.useDefaultEnvironment();
+    }
   }
 
   @Test
   public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Exception {
-    try (SharedBackendReplicaHarness harness = SharedBackendReplicaHarness.create(2);
-        Spanner spanner = createSpanner(harness)) {
+    try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
 
@@ -173,8 +187,7 @@ public void singleUseReadReroutesOnResourceExhaustedForBypassTraffic() throws Ex
 
   @Test
   public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() throws Exception {
-    try (SharedBackendReplicaHarness harness = SharedBackendReplicaHarness.create(2);
-        Spanner spanner = createSpanner(harness)) {
+    try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
 
@@ -260,8 +273,7 @@ public void singleUseReadCooldownSkipsReplicaOnNextRequestForBypassTraffic() thr
 
   @Test
   public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exception {
-    try (SharedBackendReplicaHarness harness = SharedBackendReplicaHarness.create(2);
-        Spanner spanner = createSpanner(harness)) {
+    try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
 
@@ -332,8 +344,7 @@ public void singleUseReadReroutesOnUnavailableForBypassTraffic() throws Exceptio
   @Test
   public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTraffic()
       throws Exception {
-    try (SharedBackendReplicaHarness harness = SharedBackendReplicaHarness.create(2);
-        Spanner spanner = createSpanner(harness)) {
+    try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"));
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
 
@@ -419,8 +430,7 @@ public void singleUseReadCooldownSkipsUnavailableReplicaOnNextRequestForBypassTr
   @Test
   public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTraffic()
       throws Exception {
-    try (SharedBackendReplicaHarness harness = SharedBackendReplicaHarness.create(2);
-        Spanner spanner = createSpanner(harness)) {
+    try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, multiRowReadResultSet("b", "c", "d"));
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
 
@@ -501,8 +511,7 @@ public void singleUseReadMidStreamRecvFailureWithoutRetryInfoRetriesForBypassTra
   @Test
   public void readWriteTransactionAbortedCommitUsesReadAffinityReplicaForBypassTraffic()
       throws Exception {
-    try (SharedBackendReplicaHarness harness = SharedBackendReplicaHarness.create(2);
-        Spanner spanner = createSpanner(harness)) {
+    try (Spanner spanner = createSpanner(harness)) {
       configureBackend(harness, singleRowReadResultSet("b"), /* leaderReplicaIndex= */ 1);
       DatabaseClient client = spanner.getDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, DATABASE));
 

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/SharedBackendReplicaHarness.java

@@ -97,6 +97,10 @@ synchronized void clearRequests() {
       requestIds.clear();
     }
 
+    synchronized void clearMethodErrors() {
+      methodErrors.clear();
+    }
+
     private synchronized void recordRequest(String method, AbstractMessage request) {
       requests.computeIfAbsent(method, ignored -> new ArrayList<>()).add(request);
     }
@@ -287,6 +291,18 @@ void clearRequests() {
     }
   }
 
+  void reset() {
+    backend.reset();
+    backend.removeAllExecutionTimes();
+    backend.setAbortProbability(0.0D);
+    defaultReplica.clearRequests();
+    defaultReplica.clearMethodErrors();
+    for (HookedReplicaSpannerService replica : replicas) {
+      replica.clearRequests();
+      replica.clearMethodErrors();
+    }
+  }
+
   @Override
   public void close() throws IOException {
     IOException failure = null;