unit test case added for retry mechanism

Author: AnkitCLI

src/main/java/io/cdap/plugin/gcp/gcs/ServiceAccountAccessTokenProvider.java

@@ -59,6 +59,7 @@ public class ServiceAccountAccessTokenProvider implements AccessTokenProvider {
   private static final Pattern SERVER_ERROR_PATTERN = Pattern.compile("Unexpected Error code 5\\d{2} trying to get " +
   "security access token from Compute Engine metadata for the default service account.*");
 
+  @VisibleForTesting
   @Override
   public AccessToken getAccessToken() {
       try {
@@ -100,7 +101,7 @@ static boolean isServerError(IOException e) {
     return msg != null && SERVER_ERROR_PATTERN.matcher(msg).matches();
   }
 
-  private com.google.auth.oauth2.AccessToken retrieveAccessToken() throws IOException {
+  com.google.auth.oauth2.AccessToken retrieveAccessToken() throws IOException {
     try {
       return getCredentials().getAccessToken();
     } catch (IOException e) {

src/test/java/io/cdap/plugin/gcp/gcs/ServiceAccountAccessTokenProviderTest.java

@@ -17,17 +17,23 @@
 package io.cdap.plugin.gcp.gcs;
 
 import com.google.api.client.auth.oauth2.Credential;
+import com.google.auth.oauth2.AccessToken;
 import com.google.cloud.hadoop.fs.gcs.GoogleHadoopFileSystemConfiguration;
 import com.google.cloud.hadoop.util.AccessTokenProvider;
 import com.google.cloud.hadoop.util.CredentialFactory;
 import com.google.cloud.hadoop.util.CredentialFromAccessTokenProviderClassFactory;
 import com.google.cloud.hadoop.util.HadoopCredentialConfiguration;
 import com.google.common.collect.ImmutableList;
 import io.cdap.plugin.gcp.common.GCPUtils;
+import io.cdap.plugin.gcp.common.ServerErrorException;
 import org.apache.hadoop.conf.Configuration;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.Mockito;
+
 import java.io.IOException;
+import java.time.Instant;
+import java.util.Date;
 import java.util.Map;
 
 /**
@@ -102,4 +108,36 @@ public void testIsServerErrorWith5xxErrorCode503() {
         "service account.");
     Assert.assertTrue(ServiceAccountAccessTokenProvider.isServerError(serverError));
   }
+
+  @Test(expected = ServerErrorException.class)
+  public void testRetryMechanismFailsAfterMaxRetries() throws IOException {
+    ServiceAccountAccessTokenProvider provider = Mockito.spy(new ServiceAccountAccessTokenProvider());
+    Mockito.doThrow(new ServerErrorException(503, "Unexpected Error code 503 trying to get security access token " +
+        "from Compute Engine metadata for the default service account.", null))
+      .when(provider).retrieveAccessToken();
+    provider.getAccessToken();
+  }
+
+  @Test
+  public void testRetryMechanismSucceedsAfterFewRetries() throws IOException {
+    ServiceAccountAccessTokenProvider provider = Mockito.spy(new ServiceAccountAccessTokenProvider());
+
+    // Create a valid token with future expiration
+    AccessToken validToken = new AccessToken("valid-token", Date.from(Instant.now().plusSeconds(3600)));
+
+    // Fail first 2 attempts, then succeed
+    Mockito.doThrow(new ServerErrorException(503, "Unexpected Error code 503 trying to get security access token " +
+        "from Compute Engine metadata for the default service account.", null))
+           .doThrow(new ServerErrorException(500, "Unexpected Error code 500 trying to get security access token " +
+        "from Compute Engine metadata for the default service account.", null))
+      .doReturn(validToken)
+      .when(provider).retrieveAccessToken();
+    AccessTokenProvider.AccessToken accessToken = provider.getAccessToken();
+
+    Assert.assertNotNull(accessToken);
+    Assert.assertEquals("valid-token", accessToken.getToken());
+
+    // Verify that retrieveAccessToken was called 3 times (2 failures + 1 success)
+    Mockito.verify(provider, Mockito.times(3)).retrieveAccessToken();
+  }
 }