Private 2 Protected

cloudsufi · Nov 6, 2023 · 17c2e89 · 17c2e89
1 parent ceeda27
commit 17c2e89
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 25 deletions.
diff --git a/src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryExecute.java b/src/main/java/io/cdap/plugin/gcp/bigquery/action/BigQueryExecute.java
@@ -144,10 +144,9 @@ public void run(ActionContext context) throws Exception {
             bigQuery, queryConfig, context);
   }
 
-  private void executeQueryWithExponentialBackoff(Duration initialRetryDuration, Duration maxRetryDuration,
+  protected void executeQueryWithExponentialBackoff(Duration initialRetryDuration, Duration maxRetryDuration,
                                                    int multiplier, int maxRetryCount, BigQuery bigQuery,
-                                                   QueryJobConfiguration queryConfig, ActionContext context)
-      throws ConnectException {
+                                                   QueryJobConfiguration queryConfig, ActionContext context) {
     RetryPolicy<Object> retryPolicy = RetryPolicy.builder()
             .handle(ConnectException.class)
             .withBackoff(initialRetryDuration, maxRetryDuration, multiplier)

diff --git a/src/test/java/io/cdap/plugin/gcp/bigquery/action/BigQueryExecuteTest.java b/src/test/java/io/cdap/plugin/gcp/bigquery/action/BigQueryExecuteTest.java
@@ -20,8 +20,7 @@
 import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
-import java.lang.reflect.InvocationTargetException;
-import java.lang.reflect.Method;
+import java.net.ConnectException;
 import java.time.Duration;
 
 public class BigQueryExecuteTest {
@@ -43,8 +42,6 @@ public class BigQueryExecuteTest {
   StageMetrics stageMetrics;
   @Mock
   Metrics metrics;
-  Method executeQueryWithExponentialBackoffMethod;
-
   QueryJobConfiguration queryJobConfiguration;
   BigQueryExecute.Config config;
   JobInfo jobInfo;
@@ -66,10 +63,6 @@ public void setUp() throws InterruptedException, NoSuchMethodException {
     jobId = JobId.newBuilder().setRandomJob().setLocation(config.getLocation()).build();
     jobInfo = JobInfo.newBuilder(queryJobConfiguration).setJobId(jobId).build();
     bq = new BigQueryExecute(config);
-    executeQueryWithExponentialBackoffMethod = bq.getClass()
-            .getDeclaredMethod("executeQueryWithExponentialBackoff", Duration.class, Duration.class, int.class,
-                    int.class, BigQuery.class, QueryJobConfiguration.class, ActionContext.class);
-    executeQueryWithExponentialBackoffMethod.setAccessible(true);
 
     // Mock Job Creation
     Mockito.when(bigQuery.create((JobInfo) Mockito.any())).thenReturn(queryJob);
@@ -91,28 +84,25 @@ public void setUp() throws InterruptedException, NoSuchMethodException {
 
   }
 
-  @Test(expected = java.lang.reflect.InvocationTargetException.class)
-  public void testExecuteQueryWithExponentialBackoffFailsWithNonRetryError() throws InvocationTargetException,
-          IllegalAccessException {
+  @Test(expected = java.lang.RuntimeException.class)
+  public void testExecuteQueryWithExponentialBackoffFailsWithNonRetryError() throws ConnectException {
     Mockito.when(bigQueryError.getReason()).thenReturn("accessDenied");
-    executeQueryWithExponentialBackoffMethod.invoke(bq, initialRetryDuration, maxRetryDuration, multiplier,
-            maxRetryCount, bigQuery, queryJobConfiguration, context);
+    bq.executeQueryWithExponentialBackoff(initialRetryDuration, maxRetryDuration, multiplier, maxRetryCount,
+            bigQuery, queryJobConfiguration, context);
   }
 
-  @Test(expected = java.lang.reflect.InvocationTargetException.class)
-  public void testExecuteQueryWithExponentialBackoffFailsRetryError() throws InvocationTargetException,
-          IllegalAccessException {
+  @Test(expected = java.lang.RuntimeException.class)
+  public void testExecuteQueryWithExponentialBackoffFailsRetryError() throws ConnectException {
     Mockito.when(bigQueryError.getReason()).thenReturn("jobBackendError");
-    executeQueryWithExponentialBackoffMethod.invoke(bq, initialRetryDuration, maxRetryDuration, multiplier,
-            maxRetryCount, bigQuery, queryJobConfiguration, context);
+    bq.executeQueryWithExponentialBackoff(initialRetryDuration, maxRetryDuration, multiplier, maxRetryCount,
+            bigQuery, queryJobConfiguration, context);
   }
 
   @Test
-  public void testExecuteQueryWithExponentialBackoffSuccess() throws InterruptedException,
-          InvocationTargetException, IllegalAccessException {
+  public void testExecuteQueryWithExponentialBackoffSuccess() throws InterruptedException, ConnectException {
     Mockito.when(jobStatus.getError()).thenReturn(null);
     Mockito.when(queryJob.getQueryResults()).thenReturn(queryResults);
-    executeQueryWithExponentialBackoffMethod.invoke(bq, initialRetryDuration, maxRetryDuration, multiplier,
-            maxRetryCount, bigQuery, queryJobConfiguration, context);
+    bq.executeQueryWithExponentialBackoff(initialRetryDuration, maxRetryDuration, multiplier, maxRetryCount,
+            bigQuery, queryJobConfiguration, context);
   }
 }