test(auth): add comprehensive coverage for auth management key feature

src/main/java/com/descope/client/Config.java

@@ -30,6 +30,13 @@ public class Config {
   // fail.
   private String managementKey;
 
+  // AuthManagementKey (optional, "") - used to provide a management key to use
+  // with Authentication APIs whose public access has been disabled.
+  // If empty, this value is retrieved from the DESCOPE_AUTH_MANAGEMENT_KEY
+  // environment variable instead. If neither values are set then any disabled
+  // authentication methods API calls will fail.
+  private String authManagementKey;
+
   // PublicKey (optional, "") - used to override or implicitly use a dedicated public key in order
   // to decrypt and validate the JWT tokens during ValidateSessionRequest().
   // If empty, will attempt to fetch all public keys from the specified project id.
@@ -73,4 +80,11 @@ public String initializeManagementKey() {
     }
     return this.managementKey;
   }
+
+  public String initializeAuthManagementKey() {
+    if (StringUtils.isBlank(this.authManagementKey)) {
+      this.authManagementKey = EnvironmentUtils.getAuthManagementKey();
+    }
+    return this.authManagementKey;
+  }
 }

src/main/java/com/descope/client/DescopeClient.java

@@ -49,6 +49,7 @@ public DescopeClient(Config config) throws DescopeException {
       log.debug("Provided public key is set, forcing only provided public key validation");
     }
     config.initializeManagementKey();
+    config.initializeAuthManagementKey();
     config.initializeBaseURL();
 
     Client client = getClient(config);
@@ -69,6 +70,7 @@ private static Client getClient(Config config) {
         .uri(StringUtils.isBlank(config.getDescopeBaseUrl()) ? baseUrl : config.getDescopeBaseUrl())
         .projectId(projectId)
         .managementKey(config.getManagementKey())
+        .authManagementKey(config.getAuthManagementKey())
         .headers(
             Collections.isEmpty(config.getCustomDefaultHeaders())
               ? new HashMap<>() : new HashMap<>(config.getCustomDefaultHeaders()))

src/main/java/com/descope/literals/AppConstants.java

@@ -7,6 +7,7 @@ public class AppConstants {
   public static final String PROJECT_ID_ENV_VAR = "DESCOPE_PROJECT_ID";
   public static final String PUBLIC_KEY_ENV_VAR = "DESCOPE_PUBLIC_KEY";
   public static final String MANAGEMENT_KEY_ENV_VAR = "DESCOPE_MANAGEMENT_KEY";
+  public static final String AUTH_MANAGEMENT_KEY_ENV_VAR = "DESCOPE_AUTH_MANAGEMENT_KEY";
   public static final String BASE_URL_ENV_VAR = "DESCOPE_BASE_URL";
   public static final String AUTHORIZATION_HEADER_NAME = "Authorization";
   public static final String BEARER_AUTHORIZATION_PREFIX = "Bearer ";

src/main/java/com/descope/model/client/Client.java

@@ -18,6 +18,7 @@ public class Client {
   private String uri;
   private String projectId;
   private String managementKey;
+  private String authManagementKey;
   private Map<String, String> headers;
   private SdkInfo sdkInfo;
   private Key providedKey;

src/main/java/com/descope/sdk/auth/impl/AuthenticationsBase.java

@@ -37,7 +37,11 @@ abstract class AuthenticationsBase extends SdkServicesBase implements Authentica
 
   ApiProxy getApiProxy() {
     String projectId = client.getProjectId();
+    String authManagementKey = client.getAuthManagementKey();
     if (StringUtils.isNotBlank(projectId)) {
+      if (StringUtils.isNotBlank(authManagementKey)) {
+        return ApiProxyBuilder.buildProxy(() -> String.format("Bearer %s:%s", projectId, authManagementKey), client);
+      }
       return ApiProxyBuilder.buildProxy(() -> "Bearer " + projectId, client);
     }
     return ApiProxyBuilder.buildProxy(client.getSdkInfo());
@@ -49,7 +53,13 @@ ApiProxy getApiProxy(String refreshToken) {
       return getApiProxy();
     }
 
-    String token = String.format("Bearer %s:%s", projectId, refreshToken);
+    String authManagementKey = client.getAuthManagementKey();
+    String token;
+    if (StringUtils.isNotBlank(authManagementKey)) {
+      token = String.format("Bearer %s:%s:%s", projectId, refreshToken, authManagementKey);
+    } else {
+      token = String.format("Bearer %s:%s", projectId, refreshToken);
+    }
     return ApiProxyBuilder.buildProxy(() -> token, client);
   }
 

src/main/java/com/descope/utils/EnvironmentUtils.java

@@ -1,5 +1,6 @@
 package com.descope.utils;
 
+import static com.descope.literals.AppConstants.AUTH_MANAGEMENT_KEY_ENV_VAR;
 import static com.descope.literals.AppConstants.BASE_URL_ENV_VAR;
 import static com.descope.literals.AppConstants.MANAGEMENT_KEY_ENV_VAR;
 import static com.descope.literals.AppConstants.PROJECT_ID_ENV_VAR;
@@ -27,4 +28,8 @@ public static String getPublicKey() {
   public static String getManagementKey() {
     return dotenv.get(MANAGEMENT_KEY_ENV_VAR);
   }
+
+  public static String getAuthManagementKey() {
+    return dotenv.get(AUTH_MANAGEMENT_KEY_ENV_VAR);
+  }
 }

src/test/java/com/descope/client/DescopeClientTest.java

@@ -1,5 +1,6 @@
 package com.descope.client;
 
+import static com.descope.literals.AppConstants.AUTH_MANAGEMENT_KEY_ENV_VAR;
 import static com.descope.literals.AppConstants.MANAGEMENT_KEY_ENV_VAR;
 import static com.descope.literals.AppConstants.PROJECT_ID_ENV_VAR;
 import static com.descope.literals.AppConstants.PUBLIC_KEY_ENV_VAR;
@@ -123,4 +124,62 @@ void testEmptyConfig() {
         .isInstanceOf(ServerCommonException.class)
         .hasMessage("The Config argument is invalid");
   }
+
+  @Test
+  void testAuthManagementKeyFromEnvVariable() throws Exception {
+    String expectedProjectID = "P123456789012345678901234567";
+    String expectedAuthManagementKey = "someAuthManagementKey";
+    EnvironmentVariables env =
+        new EnvironmentVariables(PROJECT_ID_ENV_VAR, expectedProjectID)
+            .and(AUTH_MANAGEMENT_KEY_ENV_VAR, expectedAuthManagementKey);
+    env.execute(
+        () -> {
+          DescopeClient descopeClient = new DescopeClient();
+          Config config = descopeClient.getConfig();
+          Assertions.assertThat(config.getProjectId()).isEqualTo(expectedProjectID);
+          Assertions.assertThat(config.getAuthManagementKey()).isEqualTo(expectedAuthManagementKey);
+        });
+  }
+
+  @Test
+  void testAuthManagementKeyFromConfig() throws Exception {
+    String expectedProjectID = "P123456789012345678901234567";
+    String expectedAuthManagementKey = "someAuthManagementKey";
+    Config config = Config.builder()
+        .projectId(expectedProjectID)
+        .authManagementKey(expectedAuthManagementKey)
+        .build();
+    DescopeClient descopeClient = new DescopeClient(config);
+    Assertions.assertThat(descopeClient.getConfig().getProjectId()).isEqualTo(expectedProjectID);
+    Assertions.assertThat(descopeClient.getConfig().getAuthManagementKey()).isEqualTo(expectedAuthManagementKey);
+  }
+
+  @Test
+  void testAuthManagementKeyAndManagementKeyTogether() throws Exception {
+    String expectedProjectID = "P123456789012345678901234567";
+    String expectedManagementKey = "someManagementKey";
+    String expectedAuthManagementKey = "someAuthManagementKey";
+    EnvironmentVariables env =
+        new EnvironmentVariables(PROJECT_ID_ENV_VAR, expectedProjectID)
+            .and(MANAGEMENT_KEY_ENV_VAR, expectedManagementKey)
+            .and(AUTH_MANAGEMENT_KEY_ENV_VAR, expectedAuthManagementKey);
+    env.execute(
+        () -> {
+          DescopeClient descopeClient = new DescopeClient();
+          Config config = descopeClient.getConfig();
+          Assertions.assertThat(config.getProjectId()).isEqualTo(expectedProjectID);
+          Assertions.assertThat(config.getManagementKey()).isEqualTo(expectedManagementKey);
+          Assertions.assertThat(config.getAuthManagementKey()).isEqualTo(expectedAuthManagementKey);
+        });
+  }
+
+  @Test
+  void testAuthManagementKeyConfigMethod() throws Exception {
+    String expectedProjectID = "P123456789012345678901234567";
+    Config config = Config.builder()
+        .projectId(expectedProjectID)
+        .build();
+    DescopeClient descopeClient = new DescopeClient(config);
+    Assertions.assertThat(descopeClient.getConfig().getProjectId()).isEqualTo(expectedProjectID);
+  }
 }

src/test/java/com/descope/sdk/TestUtils.java

@@ -131,6 +131,7 @@ public static Client getClient() {
         .uri(baseUrl)
         .projectId(EnvironmentUtils.getProjectId())
         .managementKey(EnvironmentUtils.getManagementKey())
+        .authManagementKey(EnvironmentUtils.getAuthManagementKey())
         .sdkInfo(getSdkInfo())
         .build();
   }

src/test/java/com/descope/sdk/auth/impl/AuthenticationsBaseTest.java

@@ -0,0 +1,147 @@
+package com.descope.sdk.auth.impl;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.descope.model.client.Client;
+import com.descope.model.client.SdkInfo;
+import com.descope.proxy.ApiProxy;
+import com.descope.proxy.impl.ApiProxyBuilder;
+import java.util.function.Supplier;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.MockedStatic;
+
+public class AuthenticationsBaseTest {
+
+  @Test
+  void testGetApiProxyWithAuthManagementKey() {
+    String projectId = "P123456789012345678901234567";
+    String authManagementKey = "auth-mgmt-key-123";
+
+    Client client = Client.builder()
+        .projectId(projectId)
+        .authManagementKey(authManagementKey)
+        .sdkInfo(SdkInfo.builder().name("test").build())
+        .build();
+
+    ApiProxy mockProxy = mock(ApiProxy.class);
+    ArgumentCaptor<Supplier<String>> authHeaderCaptor = ArgumentCaptor.forClass(Supplier.class);
+
+    try (MockedStatic<ApiProxyBuilder> mockedBuilder = mockStatic(ApiProxyBuilder.class)) {
+      mockedBuilder.when(() -> ApiProxyBuilder.buildProxy(authHeaderCaptor.capture(), any(Client.class)))
+          .thenReturn(mockProxy);
+
+      OTPServiceImpl otpService = new OTPServiceImpl(client);
+      otpService.getApiProxy();
+
+      String authHeader = authHeaderCaptor.getValue().get();
+      assertThat(authHeader).isEqualTo("Bearer " + projectId + ":" + authManagementKey);
+    }
+  }
+
+  @Test
+  void testGetApiProxyWithoutAuthManagementKey() {
+    String projectId = "P123456789012345678901234567";
+
+    Client client = Client.builder()
+        .projectId(projectId)
+        .sdkInfo(SdkInfo.builder().name("test").build())
+        .build();
+
+    ApiProxy mockProxy = mock(ApiProxy.class);
+    ArgumentCaptor<Supplier<String>> authHeaderCaptor = ArgumentCaptor.forClass(Supplier.class);
+
+    try (MockedStatic<ApiProxyBuilder> mockedBuilder = mockStatic(ApiProxyBuilder.class)) {
+      mockedBuilder.when(() -> ApiProxyBuilder.buildProxy(authHeaderCaptor.capture(), any(Client.class)))
+          .thenReturn(mockProxy);
+
+      OTPServiceImpl otpService = new OTPServiceImpl(client);
+      otpService.getApiProxy();
+
+      String authHeader = authHeaderCaptor.getValue().get();
+      assertThat(authHeader).isEqualTo("Bearer " + projectId);
+    }
+  }
+
+  @Test
+  void testGetApiProxyWithRefreshTokenAndAuthManagementKey() {
+    String projectId = "P123456789012345678901234567";
+    String authManagementKey = "auth-mgmt-key-123";
+    String refreshToken = "refresh-token-456";
+
+    Client client = Client.builder()
+        .projectId(projectId)
+        .authManagementKey(authManagementKey)
+        .sdkInfo(SdkInfo.builder().name("test").build())
+        .build();
+
+    ApiProxy mockProxy = mock(ApiProxy.class);
+    ArgumentCaptor<Supplier<String>> authHeaderCaptor = ArgumentCaptor.forClass(Supplier.class);
+
+    try (MockedStatic<ApiProxyBuilder> mockedBuilder = mockStatic(ApiProxyBuilder.class)) {
+      mockedBuilder.when(() -> ApiProxyBuilder.buildProxy(authHeaderCaptor.capture(), any(Client.class)))
+          .thenReturn(mockProxy);
+
+      OTPServiceImpl otpService = new OTPServiceImpl(client);
+      otpService.getApiProxy(refreshToken);
+
+      String authHeader = authHeaderCaptor.getValue().get();
+      assertThat(authHeader).isEqualTo("Bearer " + projectId + ":" + refreshToken + ":" + authManagementKey);
+    }
+  }
+
+  @Test
+  void testGetApiProxyWithRefreshTokenWithoutAuthManagementKey() {
+    String projectId = "P123456789012345678901234567";
+    String refreshToken = "refresh-token-456";
+
+    Client client = Client.builder()
+        .projectId(projectId)
+        .sdkInfo(SdkInfo.builder().name("test").build())
+        .build();
+
+    ApiProxy mockProxy = mock(ApiProxy.class);
+    ArgumentCaptor<Supplier<String>> authHeaderCaptor = ArgumentCaptor.forClass(Supplier.class);
+
+    try (MockedStatic<ApiProxyBuilder> mockedBuilder = mockStatic(ApiProxyBuilder.class)) {
+      mockedBuilder.when(() -> ApiProxyBuilder.buildProxy(authHeaderCaptor.capture(), any(Client.class)))
+          .thenReturn(mockProxy);
+
+      OTPServiceImpl otpService = new OTPServiceImpl(client);
+      otpService.getApiProxy(refreshToken);
+
+      String authHeader = authHeaderCaptor.getValue().get();
+      assertThat(authHeader).isEqualTo("Bearer " + projectId + ":" + refreshToken);
+    }
+  }
+
+  @Test
+  void testGetApiProxyWithEmptyAuthManagementKey() {
+    String projectId = "P123456789012345678901234567";
+
+    Client client = Client.builder()
+        .projectId(projectId)
+        .authManagementKey("")
+        .sdkInfo(SdkInfo.builder().name("test").build())
+        .build();
+
+    ApiProxy mockProxy = mock(ApiProxy.class);
+    ArgumentCaptor<Supplier<String>> authHeaderCaptor = ArgumentCaptor.forClass(Supplier.class);
+
+    try (MockedStatic<ApiProxyBuilder> mockedBuilder = mockStatic(ApiProxyBuilder.class)) {
+      mockedBuilder.when(() -> ApiProxyBuilder.buildProxy(authHeaderCaptor.capture(), any(Client.class)))
+          .thenReturn(mockProxy);
+
+      OTPServiceImpl otpService = new OTPServiceImpl(client);
+      otpService.getApiProxy();
+
+      String authHeader = authHeaderCaptor.getValue().get();
+      assertThat(authHeader).isEqualTo("Bearer " + projectId);
+    }
+  }
+}