Making RetryExecutor accept RetrySettings when creating it

jainsahab · jainsahab · commit 1dfafb7f8adc · 2022-09-14T10:18:32.000+05:30
diff --git a/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java b/google-cloud-datastore/src/main/java/com/google/cloud/datastore/RetryExecutor.java
@@ -15,17 +15,27 @@
  */
 package com.google.cloud.datastore;
 
-import com.google.api.core.ApiClock;
+import static com.google.cloud.RetryHelper.runWithRetries;
+
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.RetryHelper;
 import java.util.concurrent.Callable;
 
 public class RetryExecutor {
 
-  public <T> T execute(Callable<T> block, RetrySettings retrySettings,
-      ResultRetryAlgorithm<?> resultRetryAlgorithm, ApiClock clock) {
-    return RetryHelper.runWithRetries(block, retrySettings, resultRetryAlgorithm, clock);
+  private final RetrySettings retrySettings;
+  private final DatastoreOptions datastoreOptions;
+
+  public RetryExecutor(RetrySettings retrySettings,
+      DatastoreOptions datastoreOptions) {
+    this.retrySettings = retrySettings;
+    this.datastoreOptions = datastoreOptions;
+  }
+
+  public <T> T execute(Callable<T> block, ResultRetryAlgorithm<?> resultRetryAlgorithm) {
+    return runWithRetries(block, retrySettings, resultRetryAlgorithm,
+        this.datastoreOptions.getClock());
   }
 
 }
diff --git a/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java b/google-cloud-datastore/src/test/java/com/google/cloud/datastore/RetryExecutorTest.java
@@ -18,25 +18,27 @@
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
 
-import com.google.api.core.NanoClock;
 import com.google.api.gax.retrying.BasicResultRetryAlgorithm;
 import com.google.api.gax.retrying.RetrySettings;
 import java.util.concurrent.Callable;
 import org.junit.Test;
 
 public class RetryExecutorTest {
 
-  private RetryExecutor retryExecutor = new RetryExecutor();
+  private static final int MAX_ATTEMPTS = 4;
+
+  private final RetrySettings retrySettings = RetrySettings.newBuilder()
+      .setMaxAttempts(MAX_ATTEMPTS)
+      .build();
+  private final DatastoreOptions datastoreOptions = DatastoreOptions.getDefaultInstance();
+  private final RetryExecutor retryExecutor = new RetryExecutor(retrySettings, datastoreOptions);
+
 
   @Test
   public void shouldRetryWhenErrorOccurred() {
-    RetrySettings retrySettings = RetrySettings.newBuilder()
-        .setMaxAttempts(4)
-        .build();
     CustomErrorCallable callable = new CustomErrorCallable(3);
 
-    retryExecutor.execute(callable, retrySettings, new BasicResultRetryAlgorithm<>(),
-        NanoClock.getDefaultClock());
+    retryExecutor.execute(callable, new BasicResultRetryAlgorithm<>());
 
     assertThat(callable.invokedTimes(), equalTo(4));
   }
