Skip to content

Commit d503043

Browse files
kunalmnnitkunalmnnit
committed
Add UTs for AbfsConfiguration
1 parent 992ac17 commit d503043

1 file changed

Lines changed: 210 additions & 0 deletions

File tree

‎hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/TestAccountConfiguration.java‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import org.apache.hadoop.fs.azurebfs.oauth2.RefreshTokenBasedTokenProvider;
3535
import org.apache.hadoop.fs.azurebfs.oauth2.UserPasswordTokenProvider;
3636
import org.apache.hadoop.fs.azurebfs.oauth2.WorkloadIdentityTokenProvider;
37+
import org.apache.hadoop.fs.azurebfs.oauth2.ClientAssertionProvider;
3738
import org.apache.hadoop.fs.azurebfs.services.AuthType;
3839
import org.apache.hadoop.test.GenericTestUtils;
3940
import org.apache.hadoop.test.LambdaTestUtils;
@@ -49,6 +50,7 @@
4950
import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_REFRESH_TOKEN;
5051
import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_USER_NAME;
5152
import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_USER_PASSWORD;
53+
import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE;
5254
import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_TOKEN_PROVIDER_TYPE_PROPERTY_NAME;
5355
import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_SAS_TOKEN_PROVIDER_TYPE;
5456
import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -79,6 +81,8 @@ public class TestAccountConfiguration {
7981
private static final String TEST_USER_PASSWORD = "userPassword";
8082
private static final String TEST_MSI_TENANT = "msiTenant";
8183
private static final String TEST_REFRESH_TOKEN = "refreshToken";
84+
private static final String TEST_CUSTOM_CLIENT_ASSERTION_PROVIDER = "org.apache.hadoop.fs.azurebfs.TestAccountConfiguration$MockClientAssertionProvider";
85+
private static final String TEST_TOKEN_FILE = "/tmp/test-token-file";
8286

8387
private static final List<String> CLIENT_CREDENTIAL_OAUTH_CONFIG_KEYS =
8488
Collections.unmodifiableList(Arrays.asList(
@@ -605,6 +609,212 @@ private void unsetAuthConfig(AbfsConfiguration abfsConf, boolean isAccountSettin
605609
abfsConf.unset(FS_AZURE_ACCOUNT_OAUTH_USER_PASSWORD + accountNameSuffix);
606610
abfsConf.unset(FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT + accountNameSuffix);
607611
abfsConf.unset(FS_AZURE_ACCOUNT_OAUTH_REFRESH_TOKEN + accountNameSuffix);
612+
abfsConf.unset(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE + accountNameSuffix);
613+
}
614+
615+
/**
616+
* Mock implementation of ClientAssertionProvider for testing
617+
*/
618+
public static class MockClientAssertionProvider implements ClientAssertionProvider {
619+
@Override
620+
public void initialize(Configuration configuration, String accountName) throws IOException {
621+
// Mock implementation
622+
}
623+
624+
@Override
625+
public String getClientAssertion() throws IOException {
626+
return "mock-jwt-token";
627+
}
628+
}
629+
630+
/**
631+
* Test that WorkloadIdentityTokenProvider can be configured with custom ClientAssertionProvider
632+
*/
633+
@Test
634+
public void testWorkloadIdentityTokenProviderWithCustomClientAssertionProvider() throws Exception {
635+
final String accountName = "account";
636+
final Configuration conf = new Configuration();
637+
final AbfsConfiguration abfsConf = new AbfsConfiguration(conf, accountName);
638+
639+
final String accountNameSuffix = "." + abfsConf.getAccountName();
640+
641+
// Set up OAuth with WorkloadIdentityTokenProvider
642+
abfsConf.set(FS_AZURE_ACCOUNT_AUTH_TYPE_PROPERTY_NAME + accountNameSuffix, AuthType.OAuth.toString());
643+
abfsConf.set(FS_AZURE_ACCOUNT_TOKEN_PROVIDER_TYPE_PROPERTY_NAME + accountNameSuffix,
644+
WorkloadIdentityTokenProvider.class.getName());
645+
646+
// Set required OAuth parameters
647+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT + accountNameSuffix, TEST_MSI_TENANT);
648+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ID + accountNameSuffix, TEST_CLIENT_ID);
649+
650+
// Set custom ClientAssertionProvider
651+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE + accountNameSuffix,
652+
TEST_CUSTOM_CLIENT_ASSERTION_PROVIDER);
653+
654+
AccessTokenProvider tokenProvider = abfsConf.getTokenProvider();
655+
Assertions.assertThat(tokenProvider)
656+
.describedAs("Should create WorkloadIdentityTokenProvider with custom ClientAssertionProvider")
657+
.isInstanceOf(WorkloadIdentityTokenProvider.class);
658+
659+
// Verify that the custom provider configuration was read and used
660+
String customProviderType = abfsConf.getPasswordString(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE);
661+
Assertions.assertThat(customProviderType)
662+
.describedAs("Custom provider type should be configured")
663+
.isEqualTo(TEST_CUSTOM_CLIENT_ASSERTION_PROVIDER);
664+
}
665+
666+
/**
667+
* Test that WorkloadIdentityTokenProvider falls back to file-based approach when no custom provider is configured
668+
*/
669+
@Test
670+
public void testWorkloadIdentityTokenProviderWithFileBasedFallback() throws Exception {
671+
final String accountName = "account";
672+
final Configuration conf = new Configuration();
673+
final AbfsConfiguration abfsConf = new AbfsConfiguration(conf, accountName);
674+
675+
final String accountNameSuffix = "." + abfsConf.getAccountName();
676+
677+
// Set up OAuth with WorkloadIdentityTokenProvider
678+
abfsConf.set(FS_AZURE_ACCOUNT_AUTH_TYPE_PROPERTY_NAME + accountNameSuffix, AuthType.OAuth.toString());
679+
abfsConf.set(FS_AZURE_ACCOUNT_TOKEN_PROVIDER_TYPE_PROPERTY_NAME + accountNameSuffix,
680+
WorkloadIdentityTokenProvider.class.getName());
681+
682+
// Set required OAuth parameters
683+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT + accountNameSuffix, TEST_MSI_TENANT);
684+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ID + accountNameSuffix, TEST_CLIENT_ID);
685+
686+
// Don't set custom provider - should fallback to file-based approach
687+
// abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE + accountNameSuffix, ...);
688+
689+
AccessTokenProvider tokenProvider = abfsConf.getTokenProvider();
690+
Assertions.assertThat(tokenProvider)
691+
.describedAs("Should create WorkloadIdentityTokenProvider with file-based fallback")
692+
.isInstanceOf(WorkloadIdentityTokenProvider.class);
693+
694+
// Verify that no custom provider is configured (should be null or empty)
695+
String customProviderType = abfsConf.getPasswordString(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE);
696+
Assertions.assertThat(customProviderType)
697+
.describedAs("No custom provider should be configured for file-based fallback")
698+
.isNull();
699+
}
700+
701+
/**
702+
* Test that invalid custom ClientAssertionProvider class name throws appropriate exception
703+
*/
704+
@Test
705+
public void testWorkloadIdentityTokenProviderWithInvalidCustomProvider() throws Exception {
706+
final String accountName = "account";
707+
final Configuration conf = new Configuration();
708+
final AbfsConfiguration abfsConf = new AbfsConfiguration(conf, accountName);
709+
710+
final String accountNameSuffix = "." + abfsConf.getAccountName();
711+
712+
// Set up OAuth with WorkloadIdentityTokenProvider
713+
abfsConf.set(FS_AZURE_ACCOUNT_AUTH_TYPE_PROPERTY_NAME + accountNameSuffix, AuthType.OAuth.toString());
714+
abfsConf.set(FS_AZURE_ACCOUNT_TOKEN_PROVIDER_TYPE_PROPERTY_NAME + accountNameSuffix,
715+
WorkloadIdentityTokenProvider.class.getName());
716+
717+
// Set required OAuth parameters
718+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT + accountNameSuffix, TEST_MSI_TENANT);
719+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ID + accountNameSuffix, TEST_CLIENT_ID);
720+
721+
// Set invalid custom ClientAssertionProvider class
722+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE + accountNameSuffix,
723+
"non.existent.InvalidProvider");
724+
725+
TokenAccessProviderException exception = LambdaTestUtils.intercept(
726+
TokenAccessProviderException.class,
727+
() -> abfsConf.getTokenProvider());
728+
729+
Assertions.assertThat(exception.getMessage())
730+
.describedAs("Should contain error about unable to load OAuth token provider class")
731+
.contains("Unable to load OAuth token provider class");
732+
}
733+
734+
/**
735+
* Test that empty/whitespace custom ClientAssertionProvider config falls back to file-based approach
736+
*/
737+
@Test
738+
public void testWorkloadIdentityTokenProviderWithEmptyCustomProviderConfig() throws Exception {
739+
final String accountName = "account";
740+
final Configuration conf = new Configuration();
741+
final AbfsConfiguration abfsConf = new AbfsConfiguration(conf, accountName);
742+
743+
final String accountNameSuffix = "." + abfsConf.getAccountName();
744+
745+
// Set up OAuth with WorkloadIdentityTokenProvider
746+
abfsConf.set(FS_AZURE_ACCOUNT_AUTH_TYPE_PROPERTY_NAME + accountNameSuffix, AuthType.OAuth.toString());
747+
abfsConf.set(FS_AZURE_ACCOUNT_TOKEN_PROVIDER_TYPE_PROPERTY_NAME + accountNameSuffix,
748+
WorkloadIdentityTokenProvider.class.getName());
749+
750+
// Set required OAuth parameters
751+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT + accountNameSuffix, TEST_MSI_TENANT);
752+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ID + accountNameSuffix, TEST_CLIENT_ID);
753+
754+
// Set empty custom ClientAssertionProvider - should fallback to file-based
755+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE + accountNameSuffix, " ");
756+
757+
AccessTokenProvider tokenProvider = abfsConf.getTokenProvider();
758+
Assertions.assertThat(tokenProvider)
759+
.describedAs("Should create WorkloadIdentityTokenProvider with file-based fallback when provider config is empty")
760+
.isInstanceOf(WorkloadIdentityTokenProvider.class);
761+
762+
// Verify that the empty provider configuration is read but treated as empty
763+
String customProviderType = abfsConf.getPasswordString(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE);
764+
Assertions.assertThat(customProviderType)
765+
.describedAs("Empty custom provider config should be present but whitespace-only")
766+
.isEqualTo(" ");
767+
768+
// Verify that when trimmed, it's empty (this is what triggers file-based fallback)
769+
Assertions.assertThat(customProviderType.trim())
770+
.describedAs("Trimmed custom provider config should be empty")
771+
.isEmpty();
772+
}
773+
774+
/**
775+
* Test that configuration precedence works for custom ClientAssertionProvider
776+
* (account-specific vs account-agnostic)
777+
*/
778+
@Test
779+
public void testWorkloadIdentityCustomProviderConfigPrecedence() throws Exception {
780+
final String accountName = "account";
781+
final Configuration conf = new Configuration();
782+
final AbfsConfiguration abfsConf = new AbfsConfiguration(conf, accountName);
783+
784+
final String accountNameSuffix = "." + abfsConf.getAccountName();
785+
786+
// Set up OAuth with WorkloadIdentityTokenProvider
787+
abfsConf.set(FS_AZURE_ACCOUNT_AUTH_TYPE_PROPERTY_NAME + accountNameSuffix, AuthType.OAuth.toString());
788+
abfsConf.set(FS_AZURE_ACCOUNT_TOKEN_PROVIDER_TYPE_PROPERTY_NAME + accountNameSuffix,
789+
WorkloadIdentityTokenProvider.class.getName());
790+
791+
// Set required OAuth parameters
792+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT + accountNameSuffix, TEST_MSI_TENANT);
793+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ID + accountNameSuffix, TEST_CLIENT_ID);
794+
795+
// Set account-agnostic custom provider (should be overridden by account-specific)
796+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE, "some.other.Provider");
797+
798+
// Set account-specific custom provider (should take precedence)
799+
abfsConf.set(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE + accountNameSuffix,
800+
TEST_CUSTOM_CLIENT_ASSERTION_PROVIDER);
801+
802+
AccessTokenProvider tokenProvider = abfsConf.getTokenProvider();
803+
Assertions.assertThat(tokenProvider)
804+
.describedAs("Should create WorkloadIdentityTokenProvider with account-specific custom provider taking precedence")
805+
.isInstanceOf(WorkloadIdentityTokenProvider.class);
806+
807+
// Verify that account-specific configuration takes precedence over account-agnostic
808+
String accountSpecificProvider = abfsConf.getPasswordString(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE);
809+
Assertions.assertThat(accountSpecificProvider)
810+
.describedAs("Account-specific custom provider should take precedence")
811+
.isEqualTo(TEST_CUSTOM_CLIENT_ASSERTION_PROVIDER);
812+
813+
// Verify that the account-agnostic setting exists but isn't used
814+
String accountAgnosticProvider = abfsConf.getRawConfiguration().get(FS_AZURE_ACCOUNT_OAUTH_CLIENT_ASSERTION_PROVIDER_TYPE);
815+
Assertions.assertThat(accountAgnosticProvider)
816+
.describedAs("Account-agnostic setting should exist but not be used")
817+
.isEqualTo("some.other.Provider");
608818
}
609819

610820
}

0 commit comments

Comments
 (0)