如何创建JUnit5自定义注解以添加测试前后执行逻辑？

能否在JUnit5中定义自定义注解，用于在测试方法执行前后运行特定代码？由于不同测试对应的前置逻辑需要传入不同参数，无法使用@Before和@After注解。举例来说，原本需在try-with-resources中通过TenantContextHolder.withTenant("one")实现租户上下文切换，希望改为通过@WithTenant("one")注解简化测试代码：

原写法：

@Test
void myTest() throws Exception {
    try (AutoCloseable autoCloseable = TenantContextHolder.withTenant("one")) {
        // test something...
    }
}

期望写法：

@Test
@WithTenant("one")
void myTest() throws Exception {
    // test something...
}

完全可以实现，借助JUnit5的扩展机制就能搞定，具体步骤如下：

1. 定义自定义注解@WithTenant #

先写一个带参数的注解，用来标记需要切换租户上下文的测试方法，同时关联对应的扩展类：

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;
import org.junit.jupiter.api.extension.ExtendWith;

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@ExtendWith(TenantExtension.class) // 让注解自动触发扩展逻辑
public @interface WithTenant {
    String value(); // 用来传递租户ID参数
}

2. 实现JUnit5扩展类 #

写一个扩展类，实现测试前后的钩子方法，处理租户上下文的开启和关闭：

import org.junit.jupiter.api.extension.AfterTestExecutionCallback;
import org.junit.jupiter.api.extension.BeforeTestExecutionCallback;
import org.junit.jupiter.api.extension.ExtensionContext;

public class TenantExtension implements BeforeTestExecutionCallback, AfterTestExecutionCallback {

    // 用Namespace隔离不同测试的状态，避免干扰
    private static final ExtensionContext.Namespace TENANT_NAMESPACE = 
        ExtensionContext.Namespace.create(TenantExtension.class);

    @Override
    public void beforeTestExecution(ExtensionContext context) throws Exception {
        // 获取当前测试方法上的@WithTenant注解
        WithTenant withTenant = context.getRequiredTestMethod().getAnnotation(WithTenant.class);
        if (withTenant != null) {
            String tenantId = withTenant.value();
            // 开启租户上下文，得到AutoCloseable实例
            AutoCloseable closeable = TenantContextHolder.withTenant(tenantId);
            // 把实例存到当前测试的专属Store中，方便后续关闭
            context.getStore(TENANT_NAMESPACE)
                   .put(context.getUniqueId(), closeable);
        }
    }

    @Override
    public void afterTestExecution(ExtensionContext context) throws Exception {
        // 从Store中取出之前保存的AutoCloseable并关闭
        AutoCloseable closeable = context.getStore(TENANT_NAMESPACE)
                                         .get(context.getUniqueId(), AutoCloseable.class);
        if (closeable != null) {
            closeable.close();
        }
    }
}

3. 使用自定义注解 #

现在就可以像你期望的那样写测试代码了，直接在测试方法上添加@WithTenant注解并传入租户ID：

import org.junit.jupiter.api.Test;

@Test
@WithTenant("one")
void myTest() throws Exception {
    // test something...
}

原理说明 #

• JUnit5的Extension机制允许我们在测试生命周期的关键节点插入自定义逻辑，这里用BeforeTestExecutionCallback在测试方法执行前开启租户上下文，用AfterTestExecutionCallback在测试执行完成后关闭上下文。
• 通过ExtensionContext.Store保存每个测试的AutoCloseable实例，利用Namespace和测试的唯一ID确保不同测试之间的状态互不干扰，避免并发测试时出现问题。

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