如何在SpringBootTest容器启动前添加系统变量

原创

半月无霜

发布于 2024-07-19 11:35:12

1080

发布于 2024-07-19 11:35:12

文章被收录于专栏：半月无霜

如何在SpringBootTest容器启动前添加系统变量

一、介绍

在上一篇文章中，我们讲解了SpringBootTest如何只对Mapper的方法进行测试

这种就是简单的启动一个SpringBoot容器就好，仅加载需要的Bean，去掉无用的controller和service，使得启动速度大幅度提升

但在结尾，留下了一个疑问，我使用了Mybatis-plus的mpw密码加密的，这没办法生效怎么办。

二、解决

如果你是用本地的配置文件，把这个类也加入进去就行，即SafetyEncryptProcessor.java

这个类是Mybatis-plus的解密

/*
 * Copyright (c) 2011-2023, baomidou (jobob@qq.com).
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.baomidou.mybatisplus.autoconfigure;

import com.baomidou.mybatisplus.core.toolkit.AES;
import com.baomidou.mybatisplus.core.toolkit.CollectionUtils;
import com.baomidou.mybatisplus.core.toolkit.StringUtils;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.env.EnvironmentPostProcessor;
import org.springframework.boot.env.OriginTrackedMapPropertySource;
import org.springframework.core.env.ConfigurableEnvironment;
import org.springframework.core.env.MapPropertySource;
import org.springframework.core.env.PropertySource;
import org.springframework.core.env.SimpleCommandLinePropertySource;

import java.util.HashMap;

/**
 * 安全加密处理器
 *
 * @author hubin
 * @since 2020-05-23
 */
public class SafetyEncryptProcessor implements EnvironmentPostProcessor {

    @Override
    public void postProcessEnvironment(ConfigurableEnvironment environment, SpringApplication application) {
        /**
         * 命令行中获取密钥
         */
        String mpwKey = null;
        for (PropertySource<?> ps : environment.getPropertySources()) {
            if (ps instanceof SimpleCommandLinePropertySource) {
                SimpleCommandLinePropertySource source = (SimpleCommandLinePropertySource) ps;
                mpwKey = source.getProperty("mpw.key");
                break;
            }
        }
        /**
         * 处理加密内容
         */
        if (StringUtils.isNotBlank(mpwKey)) {
            HashMap<String, Object> map = new HashMap<>();
            for (PropertySource<?> ps : environment.getPropertySources()) {
                if (ps instanceof OriginTrackedMapPropertySource) {
                    OriginTrackedMapPropertySource source = (OriginTrackedMapPropertySource) ps;
                    for (String name : source.getPropertyNames()) {
                        Object value = source.getProperty(name);
                        if (value instanceof String) {
                            String str = (String) value;
                            if (str.startsWith("mpw:")) {
                                map.put(name, AES.decrypt(str.substring(4), mpwKey));
                            }
                        }
                    }
                }
            }
            // 将解密的数据放入环境变量，并处于第一优先级上
            if (CollectionUtils.isNotEmpty(map)) {
                environment.getPropertySources().addFirst(new MapPropertySource("custom-encrypt", map));
            }
        }
    }
}

可是，万一用的是nacos，在远端配置的该怎么办，就连正常容器都没办法启动哦

解决方法如下，亲测可用，感谢大佬

mybatis-plus mpw 基于nacos-配置中心失效问题解决。_nacos导致数据库加密失效-CSDN博客

那么在测试的时候，只需要将DecryptNacosPropertyPluginConfig.java也加入到容器中就行

package com.banmoon.mapper;

import com.banmoon.business.config.DecryptNacosPropertyPluginConfig;
import com.baomidou.mybatisplus.autoconfigure.MybatisPlusAutoConfiguration;
import org.apache.ibatis.session.SqlSessionFactory;
import org.junit.runner.RunWith;
import org.mybatis.spring.annotation.MapperScan;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import javax.sql.DataSource;

/**
 * @author banmoon
 * @date 2024/07/16 11:36:41
 */
@SpringBootTest(classes = {
        PowerMockitoMapper.class,
        DataSourceAutoConfiguration.class,
        MybatisPlusAutoConfiguration.class,
        DataSource.class,
        SqlSessionFactory.class,
        DecryptNacosPropertyPluginConfig.class
})
@RunWith(SpringRunner.class)
@MapperScan("com.banmoon.mapper")
public class PowerMockitoMapperTest {
  
}

那么，这两个类加上去了，没什么问题

系统变量哪里设置，本地还好，在IDEA上可以这样添加

因为到时候测试Mapper方法会很多，如果不想这样，那该怎么做呢

三、如何添加系统变量

初始化的时候，直接手动的往系统变量里面添加一个，到时候执行到解密bean，能获取到不就行了

如此，我先试了@Before方法中添加，如下代码

package com.banmoon.mapper;

import com.banmoon.business.config.DecryptNacosPropertyPluginConfig;
import com.baomidou.mybatisplus.autoconfigure.MybatisPlusAutoConfiguration;
import org.apache.ibatis.session.SqlSessionFactory;
import org.junit.Before;
import org.junit.runner.RunWith;
import org.mybatis.spring.annotation.MapperScan;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import javax.sql.DataSource;

/**
 * @author banmoon
 * @date 2024/07/16 11:36:41
 */
@SpringBootTest(classes = {
        PowerMockitoMapper.class,
        DataSourceAutoConfiguration.class,
        MybatisPlusAutoConfiguration.class,
        DataSource.class,
        SqlSessionFactory.class,
        DecryptNacosPropertyPluginConfig.class
})
@RunWith(SpringRunner.class)
@MapperScan("com.banmoon.mapper")
public class PowerMockitoMapperTest {
  
    @Before
    public void setUp() {
        System.setProperty("mpw.key", "abcdefg");
    }
  
}

结果自然是失败了，在SpringBoot容器启动完成之后，才堪堪走到这里

顺序不对，导致解密失败，那么只能再换一种方式了

没想到还真有，@TestExecutionListeners，可以定义一个或多个监听器，来帮我扩展junit测试的一些行为

这边我们选择让其进行初始化，我又对其封装了一下

首先定义了一个注解，其中使用了@TestExecutionListeners

package com.banmoon.test.spv.annotation;

import com.banmoon.test.spv.listener.SystemPropertyTestExecutionListener;
import org.springframework.test.context.TestExecutionListeners;

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;

/**
 * @author banmoon
 * @date 2024/07/16 15:16:05
 */
@Target(ElementType.TYPE)
@Retention(RetentionPolicy.RUNTIME)
@TestExecutionListeners(listeners = SystemPropertyTestExecutionListener.class, mergeMode = TestExecutionListeners.MergeMode.MERGE_WITH_DEFAULTS)
public @interface SystemPropeerty {

    /**
     * 系统变量key
     */
    String key();

    /**
     * 系统变量value
     */
    String value();

}

再然后是上面指定的监听器，SystemPropertyTestExecutionListener.java

package com.banmoon.test.spv.listener;

import com.banmoon.test.spv.annotation.SystemPropeerty;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.support.AbstractTestExecutionListener;

/**
 * @author banmoon
 * @date 2024/07/16 14:02:42
 */
public class SystemPropertyTestExecutionListener extends AbstractTestExecutionListener {

    @Override
    public void beforeTestClass(TestContext testContext) {
        SystemPropeerty annotation = testContext.getTestClass().getAnnotation(SystemPropeerty.class);
        String key = annotation.key();
        String value = annotation.value();
        System.setProperty(key, value);
    }
}

很简单，就是从注解中读取属性，并设置到系统变量中，使用如下

package com.banmoon.mapper;

import com.banmoon.business.config.DecryptNacosPropertyPluginConfig;
import com.baomidou.mybatisplus.autoconfigure.MybatisPlusAutoConfiguration;
import com.banmoon.test.spv.annotation.SystemPropeerty;
import org.apache.ibatis.session.SqlSessionFactory;
import org.junit.runner.RunWith;
import org.mybatis.spring.annotation.MapperScan;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.junit4.SpringRunner;

import javax.sql.DataSource;

/**
 * @author banmoon
 * @date 2024/07/16 11:36:41
 */
@SpringBootTest(classes = {
        PowerMockitoMapper.class,
        DataSourceAutoConfiguration.class,
        MybatisPlusAutoConfiguration.class,
        DataSource.class,
        SqlSessionFactory.class,
        DecryptNacosPropertyPluginConfig.class
})
@RunWith(SpringRunner.class)
@MapperScan("com.banmoon.mapper")
@SystemPropeerty(key = "mpw.key", value = "abcdefg")
public class PowerMockitoMapperTest {
  
}