Java中使用JUnit/Mockito测试JWKS RSA签名校验及角色解析方法

原有方法内部直接实例化UrlJwkProvider、调用静态方法JWT.decode、直接new Date()获取当前时间，这些硬编码的依赖会导致单元测试无法直接隔离外部依赖（Azure AD公钥接口），我们可以先做少量可测试性改造，再基于JUnit 5 + Mockito编写测试用例。 #

1 前置准备 #

1.1 依赖引入 #

在pom.xml（Maven项目）中引入以下测试依赖：

<dependency>
    <groupId>org.junit.jupiter</groupId>
    <artifactId>junit-jupiter-api</artifactId>
    <version>5.9.2</version>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.mockito</groupId>
    <artifactId>mockito-core</artifactId>
    <version>4.11.0</version>
    <scope>test</scope>
</dependency>
<dependency>
    <groupId>org.mockito</groupId>
    <artifactId>mockito-junit-jupiter</artifactId>
    <version>4.11.0</version>
    <scope>test</scope>
</dependency>
<!-- 支持mock构造方法，不需要改造原有代码时可以引入 -->
<dependency>
    <groupId>org.mockito</groupId>
    <artifactId>mockito-inline</artifactId>
    <version>4.11.0</version>
    <scope>test</scope>
</dependency>

1.2 可测试性优化（可选但推荐） #

调整原有方法，把不稳定的依赖抽为可注入参数，避免mock静态/构造方法，同时保留原有重载方法兼容旧逻辑：

// 新增加载方法，把JwkProvider、当前时间作为参数传入，方便测试
public String[] getAadRoles(String token, String intendedAudience, JwkProvider provider, Date currentDateTime)
                    throws AadTokenAuthenticationFailedException {
    DecodedJWT jwt = JWT.decode(token);
    Jwk jwk = provider.get(jwt.getKeyId());
    Algorithm algorithm=Algorithm.RSA256((RSAPublicKey) jwk.getPublicKey(), null);
    algorithm.verify(jwt);
    String audience = jwt.getAudience().get(0);
    Date expiresAt = jwt.getExpiresAt();
    Claim roles = jwt.getClaim("roles");
    if (audience.contentEquals(intendedAudience) && currentDateTime.before(expiresAt)) {
        return roles!=null ? roles.asArray(String.class) : new String[] {};
    } else {
        String errorMessage = String.format("Token not valid for Audience {%s} , with Expiry Time as {%s}.", audience, expiresAt.toString());
        LOG.error(errorMessage);
        throw new AadTokenAuthenticationFailedException();
    }
}
// 保留原有方法兼容旧调用
public String[] getAadRoles(String token, String intendedAudience, String microsoftJwksUrl) throws AadTokenAuthenticationFailedException {
    return getAadRoles(token, intendedAudience, new UrlJwkProvider(new URL(microsoftJwksUrl)), new Date());
}

2 测试用例编写 #

2.1 测试前置初始化 #

提前生成测试用RSA密钥对和测试JWT，不需要依赖真实的AAD服务：

class AadRoleParserTest {
    private RSAPublicKey testPublicKey;
    private RSAPrivateKey testPrivateKey;
    private final String TEST_KEY_ID = "test-kid-123";
    private final String TEST_AUDIENCE = "test-app-id";
    private final String[] TEST_ROLES = new String[]{"Admin", "User"};

    // 每个测试用例执行前生成测试用RSA密钥对
    @BeforeEach
    void initRsaKeyPair() throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();
        testPublicKey = (RSAPublicKey) keyPair.getPublic();
        testPrivateKey = (RSAPrivateKey) keyPair.getPrivate();
    }

    // 工具方法：生成不同场景的测试JWT
    private String generateTestToken(Date expireTime, String audience, String[] roles, boolean invalidSignature) throws Exception {
        Algorithm algorithm = Algorithm.RSA256(testPublicKey, testPrivateKey);
        JWTCreator.Builder builder = JWT.create()
                .withKeyId(TEST_KEY_ID)
                .withAudience(audience)
                .withExpiresAt(expireTime);
        if (roles != null) {
            builder.withArrayClaim("roles", roles);
        }
        // 需要生成签名错误的token时，用错误的私钥签名
        if (invalidSignature) {
            KeyPair wrongKeyPair = KeyPairGenerator.getInstance("RSA").generateKeyPair();
            Algorithm wrongAlgorithm = Algorithm.RSA256(null, (RSAPrivateKey) wrongKeyPair.getPrivate());
            return builder.sign(wrongAlgorithm);
        }
        return builder.sign(algorithm);
    }
}

2.2 各场景测试用例 #

// 测试场景1：合法token，返回正确角色
@Test
void getAadRoles_ValidToken_ReturnCorrectRoles() throws Exception {
    // 构造测试数据
    Date expireTime = Date.from(Instant.now().plus(Duration.ofHours(1)));
    String validToken = generateTestToken(expireTime, TEST_AUDIENCE, TEST_ROLES, false);
    Date currentTime = Date.from(Instant.now());
    // Mock依赖
    JwkProvider mockProvider = Mockito.mock(JwkProvider.class);
    Jwk mockJwk = Mockito.mock(Jwk.class);
    Mockito.when(mockProvider.get(TEST_KEY_ID)).thenReturn(mockJwk);
    Mockito.when(mockJwk.getPublicKey()).thenReturn(testPublicKey);
    // 执行测试
    String[] result = getAadRoles(validToken, TEST_AUDIENCE, mockProvider, currentTime);
    // 断言结果
    Assertions.assertArrayEquals(TEST_ROLES, result);
}

// 测试场景2：签名校验失败，抛出认证异常
@Test
void getAadRoles_InvalidSignature_ThrowAuthException() throws Exception {
    Date expireTime = Date.from(Instant.now().plus(Duration.ofHours(1)));
    String invalidToken = generateTestToken(expireTime, TEST_AUDIENCE, TEST_ROLES, true);
    Date currentTime = Date.from(Instant.now());
    JwkProvider mockProvider = Mockito.mock(JwkProvider.class);
    Jwk mockJwk = Mockito.mock(Jwk.class);
    Mockito.when(mockProvider.get(TEST_KEY_ID)).thenReturn(mockJwk);
    Mockito.when(mockJwk.getPublicKey()).thenReturn(testPublicKey);
    // 断言抛出指定异常
    Assertions.assertThrows(AadTokenAuthenticationFailedException.class, () -> {
        getAadRoles(invalidToken, TEST_AUDIENCE, mockProvider, currentTime);
    });
}

// 测试场景3：受众不匹配，抛出认证异常
@Test
void getAadRoles_WrongAudience_ThrowAuthException() throws Exception {
    Date expireTime = Date.from(Instant.now().plus(Duration.ofHours(1)));
    String token = generateTestToken(expireTime, "wrong-audience", TEST_ROLES, false);
    Date currentTime = Date.from(Instant.now());
    JwkProvider mockProvider = Mockito.mock(JwkProvider.class);
    Jwk mockJwk = Mockito.mock(Jwk.class);
    Mockito.when(mockProvider.get(TEST_KEY_ID)).thenReturn(mockJwk);
    Mockito.when(mockJwk.getPublicKey()).thenReturn(testPublicKey);
    Assertions.assertThrows(AadTokenAuthenticationFailedException.class, () -> {
        getAadRoles(token, TEST_AUDIENCE, mockProvider, currentTime);
    });
}

// 测试场景4：token过期，抛出认证异常
@Test
void getAadRoles_TokenExpired_ThrowAuthException() throws Exception {
    Date expireTime = Date.from(Instant.now().minus(Duration.ofHours(1)));
    String token = generateTestToken(expireTime, TEST_AUDIENCE, TEST_ROLES, false);
    Date currentTime = Date.from(Instant.now());
    JwkProvider mockProvider = Mockito.mock(JwkProvider.class);
    Jwk mockJwk = Mockito.mock(Jwk.class);
    Mockito.when(mockProvider.get(TEST_KEY_ID)).thenReturn(mockJwk);
    Mockito.when(mockJwk.getPublicKey()).thenReturn(testPublicKey);
    Assertions.assertThrows(AadTokenAuthenticationFailedException.class, () -> {
        getAadRoles(token, TEST_AUDIENCE, mockProvider, currentTime);
    });
}

// 测试场景5：token无roles字段，返回空数组
@Test
void getAadRoles_NoRolesClaim_ReturnEmptyArray() throws Exception {
    Date expireTime = Date.from(Instant.now().plus(Duration.ofHours(1)));
    String token = generateTestToken(expireTime, TEST_AUDIENCE, null, false);
    Date currentTime = Date.from(Instant.now());
    JwkProvider mockProvider = Mockito.mock(JwkProvider.class);
    Jwk mockJwk = Mockito.mock(Jwk.class);
    Mockito.when(mockProvider.get(TEST_KEY_ID)).thenReturn(mockJwk);
    Mockito.when(mockJwk.getPublicKey()).thenReturn(testPublicKey);
    String[] result = getAadRoles(token, TEST_AUDIENCE, mockProvider, currentTime);
    Assertions.assertArrayEquals(new String[]{}, result);
}

如果不想改造原有方法，可以用Mockito的构造方法mock能力直接mockUrlJwkProvider的实例化过程，不需要修改业务代码即可完成测试。

内容的提问来源于stack exchange，提问作者codeforfun