diff --git a/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/services/TestAbfsClient.java b/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/services/TestAbfsClient.java index ce9c0325a4..0fd65fb0a6 100644 --- a/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/services/TestAbfsClient.java +++ b/hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/services/TestAbfsClient.java @@ -246,21 +246,29 @@ public static AbfsClient createTestClientFromCurrentContext( AbfsClient baseAbfsClientInstance, AbfsConfiguration abfsConfig) throws AzureBlobFileSystemException { - AbfsPerfTracker tracker = new AbfsPerfTracker("test", - abfsConfig.getAccountName(), - abfsConfig); + AuthType currentAuthType = abfsConfig.getAuthType( + abfsConfig.getAccountName()); - // Create test AbfsClient - AbfsClient testClient = new AbfsClient( - baseAbfsClientInstance.getBaseUrl(), - new SharedKeyCredentials(abfsConfig.getAccountName().substring(0, - abfsConfig.getAccountName().indexOf(DOT)), - abfsConfig.getStorageAccountKey()), - abfsConfig, - new ExponentialRetryPolicy(abfsConfig.getMaxIoRetries()), - abfsConfig.getTokenProvider(), - tracker); + AbfsPerfTracker tracker = new AbfsPerfTracker("test", + abfsConfig.getAccountName(), + abfsConfig); - return testClient; - } + // Create test AbfsClient + AbfsClient testClient = new AbfsClient( + baseAbfsClientInstance.getBaseUrl(), + (currentAuthType == AuthType.SharedKey + ? new SharedKeyCredentials( + abfsConfig.getAccountName().substring(0, + abfsConfig.getAccountName().indexOf(DOT)), + abfsConfig.getStorageAccountKey()) + : null), + abfsConfig, + new ExponentialRetryPolicy(abfsConfig.getMaxIoRetries()), + (currentAuthType == AuthType.OAuth + ? abfsConfig.getTokenProvider() + : null), + tracker); + + return testClient; + } }