原因

最近在编写单测时，发现使用 Mock 工具预定义 Service 中方法的行为特别难用，而且无法精细化的实现自定义的行为，因此想要在 Spring 容器运行过程中使用自定义 Mock 对象，该对象能够代替实际的 Bean 的给定方法。

方案

创建一个 Mock 注解，并且在 Spring 容器注册完所有的 Bean 之后，解析 classpath 下所有引入该 Mock 注解的类，使用 Mock 注解标记的 Bean 替换注解中指定名称的 Bean。这种方式类似于 mybatis-spring 动态解析 @Mapper 注解的方法（MapperScannerRegistrar 实现了@Mapper 注解的扫描），但是不一样的是 mybatis-spring 使用工厂类替换接口类，而我们是用 Mock 的 Bean 替换实际的 Bean。

实现

创建 Mock 注解

/**
 * 为指定的 Bean 创建 Mock 对象，需要继承原始 Bean
 */
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.TYPE)
public @interface FakeBeanFor {
    String value(); // 需要替换的 Bean 的名称
}

在 Spring 容器注册完所有的 Bean 后，解析 classpath 下引入 @FakeBeanFor 注解的类，使用 @FakeBeanFor 注解标记的 Bean 替换 value 中指定名称的 Bean。

/**
 * 从当前 classpath 读取 @FakeBeanFor 注解的类，并替换指定名称的 bean
 */
@Slf4j
@Configuration
@ConditionalOnExpression("${unitcases.enable.fake:true}")
// 通过 BeanDefinitionRegistryPostProcessor.postProcessBeanDefinitionRegistry 可以将 Bean 动态注入容器
// 通过 BeanFactoryAware 可以自动注入 BeanFactory
public class FakeBeanConfiguration implements BeanDefinitionRegistryPostProcessor, BeanFactoryAware {

    private BeanFactory beanFactory;

    @Override
    public void postProcessBeanDefinitionRegistry(BeanDefinitionRegistry registry) throws BeansException {
        log.debug("searching for classes annotated with @FakeBeanFor");

        // 自定义 Scanner 扫描 classpath 下的指定注解
        ClassPathFakeAnnotationScanner scanner = new ClassPathFakeAnnotationScanner(registry);
        try {
            List<String> packages = AutoConfigurationPackages.get(this.beanFactory); // 获取包路径
            if (log.isDebugEnabled()) {
                for (String pkg : packages) {
                    log.debug("Using auto-configuration base package: {}", pkg);
                }
            }
            scanner.doScan(StringUtils.toStringArray(packages)); // 扫描所有加载的包
        } catch (IllegalStateException ex) {
            log.debug("could not determine auto-configuration package, automatic fake scanning disabled.", ex);
        }
    }

    @Override
    public void postProcessBeanFactory(ConfigurableListableBeanFactory factory) throws BeansException {
        // empty
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = beanFactory;
    }

    private static class ClassPathFakeAnnotationScanner extends ClassPathBeanDefinitionScanner {

        ClassPathFakeAnnotationScanner(BeanDefinitionRegistry registry) {
            super(registry, false);
            // 设置过滤器。仅扫描 @FakeBeanFor
            addIncludeFilter(new AnnotationTypeFilter(FakeBeanFor.class));
        }

        @Override
        public Set<BeanDefinitionHolder> doScan(String... basePackages) {
            List<String> fakeClassNames = new ArrayList<>();
            // 扫描全部 package 下 annotationClass 指定的 Bean
            Set<BeanDefinitionHolder> beanDefinitions = super.doScan(basePackages);

            GenericBeanDefinition definition;
            for (BeanDefinitionHolder holder : beanDefinitions) {
                definition = (GenericBeanDefinition) holder.getBeanDefinition();

                // 获取类名，并创建 Class 对象
                String className = definition.getBeanClassName();
                Class<?> clazz = classNameToClass(className);

                // 解析注解上的 value
                FakeBeanFor annotation = clazz.getAnnotation(FakeBeanFor.class);
                if (annotation == null || StringUtils.isEmpty(annotation.value())) {
                    continue;
                }

                // 使用当前加载的 @FakeBeanFor 指定的 Bean 替换 value 里指定名称的 Bean
                if (getRegistry().containsBeanDefinition(annotation.value())) {
                    getRegistry().removeBeanDefinition(annotation.value());
                    getRegistry().registerBeanDefinition(annotation.value(), definition);
                    fakeClassNames.add(clazz.getName());
                }
            }
            log.info("found fake beans: " + fakeClassNames);

            return beanDefinitions;
        }

        // 反射通过 class 名称获取 Class 对象
        private Class<?> classNameToClass(String className) {
            try {
                return Class.forName(className);
            } catch (ClassNotFoundException e) {
                log.error("create instance failed.", e);
            }
            return null;
        }
    }
}

有点儿不一样的是这是一个配置类，将它放置到 Spring 的自动扫描路径上，就可以自动扫描 classpath 下 @FakeBeanFor 指定的类，并将其加载为 BeanDefinition。在 FakeBeanConfiguration 上还配置了 ConditionalOnExpression，这样就可以只在单测环境下的 application.properties 文件中设置指定条件使得该 Configuration 生效。

注意：这里 unitcases.enable.fake:true 默认开启了替换，如果想要默认关闭则需要设置 unitcases.enable.fake:false，并且在单测环境的 application.properties 文件设置 unitcases.enable.fake=true。

举例

假设在容器中定义如下 Service：

@Service
public class HelloService {
    public void sayHello() {
        System.out.println("hello real world!");
    }
}

在单测环境下希望能够改变它的行为，但是又不想修改这个类本身，则可以使用 @FakeBeanFor 注解：

@FakeBeanFor("helloService")
public class FakeHelloService extends HelloService {
    @Override
    public void sayHello() {
        System.out.println("hello fake world!");
    }
}

通过继承实际的 Service，并覆盖 Service 的原始方法，修改其行为。在单测中可以这样使用：

@SpringBootTest
@RunWith(SpringRunner.class)
public class FakeHelloServiceTest {

    @Autowired
    private HelloService helloService;
    
    @Test
    public void testSayHello() {
        helloService.sayHello(); // 输出：“hello fake world!”
    }
}

总结

通过自定义的 Mock 对象动态替换实际的 Bean 可以实现单测环境下比较难以使用 Mock 框架实现的功能，如将原本的异步调用逻辑修改为同步调用，避免单测完成时，异步调用还未执行完成的场景。

---

欢迎关注我的公众号：我的搬砖日记，我会定时分享自己的学习历程。