动态替换Spring容器中的Bean

原因

最近在编写单测时，发现使用 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 框架实现的功能，如将原本的异步调用逻辑修改为同步调用，避免单测完成时，异步调用还未执行完成的场景。

---

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