From 17c2e895001eb756dddbb746f7d405563c4a4e60 Mon Sep 17 00:00:00 2001
From: psainics <pushpender.saini@cloudsufi.com>
Date: Tue, 31 Oct 2023 10:22:43 +0530
Subject: [PATCH] Private 2 Protected

---
 .../gcp/bigquery/action/BigQueryExecute.java  |  5 ++-
 .../bigquery/action/BigQueryExecuteTest.java  | 34 +++++++------------
 2 files changed, 14 insertions(+), 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
index 29d5608411..c1edbe7186 100644
--- 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
index 5af0e2914b..fdbb913319 100644
--- 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);
   }
 }