单元测试加上@Transactional就能实现回滚【原理】

---

前言

有时候我们开发需求的过程中写单测，会往数据库里插入数据，但是这样的数据其实是没有意义的，要么测试数据中含有唯一主键，每次run之前都要改单测里的参数，要么，有时候公司上线的流水线会跑一遍单测，给数据库多次插入无意义的数据。那么我们就需要让单测在run完之后可以自动回滚。

一、问题

如果保证单测插入到数据库的数据回滚？

二、回答

最简单的方式就是在方法上加上@Transactional注解了

三、源码分析

1.TestContextManager

{@code TestContextManager} is the main entry point into the <em>Spring
 * TestContext Framework</em>.
  TestContextManager是Spring测试框架的主要进入点
 *
 * <p>Specifically, a {@code TestContextManager} is responsible for managing a
 * single {@link TestContext} and signaling events to all registered
 * {@link TestExecutionListener TestExecutionListeners} at the following test
 <p>具体来说，{@code TestContextManager} 负责管理单个 {@link TestContext} 并在以下测试中向所有注册的 {@link TestExecutionListener TestExecutionListeners} 发送信号事件

根据上面的Java doc，我们便知道了两个重要的概念

TestContext 和 TestExecutionListener

剩下的就是listener和testContext的交互了

2.TestContext

{@code TestContext} encapsulates the context in which a test is executed,
 * agnostic of the actual testing framework in use.
 {@code TestContext} 封装了执行测试的上下文，与使用的实际测试框架无关。

可以来看看默认实现DefaultTestContext

public class DefaultTestContext implements TestContext {

	private final Map<String, Object> attributes = new ConcurrentHashMap<>(4);

	private final CacheAwareContextLoaderDelegate cacheAwareContextLoaderDelegate;

	private final MergedContextConfiguration mergedContextConfiguration;

	private final Class<?> testClass;

	@Nullable
	private volatile Object testInstance;

	@Nullable
	private volatile Method testMethod;

	@Nullable
	private volatile Throwable testException;
}

对于我们一般的方法，最重要的其实就是testClass和testMethod

3.TestExecutionListener

3.1 AbstractTestContextBootstrapper#getTestExecutionListeners

testExecutionListeners是通过SpringFactory从spring.factories文件里的配置加载进来的
其中包含了我们本文关注的TransactionalTestExecutionListener

对于从spring.factories中加载配置类不熟悉的同学，可以参考这篇博客

3.2 钩子函数

其实钩子函数可以理解为静态AOP，也可以理解为spring框架开放给用户在某个时间点实现某些功能的扩展点。在spring源码里特别常见，譬如spring里的各种postProcessor的接口实现类，里面其实没有方法逻辑，用户可以通过自定义这些接口来在spring的bean的生命周期内对bean的实现进行个性化。
言归正传，咱们还是来看TestExecutionListener的钩子函数

/**
	 * Pre-processes a test class before execution of all tests within
	 * the class.
	 * 
This method should be called immediately before framework-specific
	 * before class lifecycle callbacks.
	 * 
The default implementation is empty. Can be overridden by
	 * concrete classes as necessary.
	 * @param testContext the test context for the test; never {@code null}
	 * @throws Exception allows any exception to propagate
	 * @since 3.0
	 */
	default void beforeTestClass(TestContext testContext) throws Exception {
		/* no-op */
	}

	/**
	 * Prepares the {@link Object test instance} of the supplied
	 * {@link TestContext test context}, for example by injecting dependencies.
	 * 
This method should be called immediately after instantiation of the test
	 * instance but prior to any framework-specific lifecycle callbacks.
	 * 
The default implementation is empty. Can be overridden by
	 * concrete classes as necessary.
	 * @param testContext the test context for the test; never {@code null}
	 * @throws Exception allows any exception to propagate
	 */
	default void prepareTestInstance(TestContext testContext) throws Exception {
		/* no-op */
	}

	/**
	 * Pre-processes a test before execution of before
	 * lifecycle callbacks of the underlying test framework — for example,
	 * by setting up test fixtures.
	 * 
This method must be called immediately prior to
	 * framework-specific before lifecycle callbacks. For historical
	 * reasons, this method is named {@code beforeTestMethod}. Since the
	 * introduction of {@link #beforeTestExecution}, a more suitable name for
	 * this method might be something like {@code beforeTestSetUp} or
	 * {@code beforeEach}; however, it is unfortunately impossible to rename
	 * this method due to backward compatibility concerns.
	 * 
The default implementation is empty. Can be overridden by
	 * concrete classes as necessary.
	 * @param testContext the test context in which the test method will be
	 * executed; never {@code null}
	 * @throws Exception allows any exception to propagate
	 * @see #afterTestMethod
	 * @see #beforeTestExecution
	 * @see #afterTestExecution
	 */
	default void beforeTestMethod(TestContext testContext) throws Exception {
		/* no-op */
	}

	/**
	 * Pre-processes a test immediately before execution of the
	 * {@link java.lang.reflect.Method test method} in the supplied
	 * {@link TestContext test context} — for example, for timing
	 * or logging purposes.
	 * 
This method must be called after framework-specific
	 * before lifecycle callbacks.
	 * 
The default implementation is empty. Can be overridden by
	 * concrete classes as necessary.
	 * @param testContext the test context in which the test method will be
	 * executed; never {@code null}
	 * @throws Exception allows any exception to propagate
	 * @since 5.0
	 * @see #beforeTestMethod
	 * @see #afterTestMethod
	 * @see #afterTestExecution
	 */
	default void beforeTestExecution(TestContext testContext) throws Exception {
		/* no-op */
	}

	/**
	 * Post-processes a test immediately after execution of the
	 * {@link java.lang.reflect.Method test method} in the supplied
	 * {@link TestContext test context} — for example, for timing
	 * or logging purposes.
	 * 
This method must be called before framework-specific
	 * after lifecycle callbacks.
	 * 
The default implementation is empty. Can be overridden by
	 * concrete classes as necessary.
	 * @param testContext the test context in which the test method will be
	 * executed; never {@code null}
	 * @throws Exception allows any exception to propagate
	 * @since 5.0
	 * @see #beforeTestMethod
	 * @see #afterTestMethod
	 * @see #beforeTestExecution
	 */
	default void afterTestExecution(TestContext testContext) throws Exception {
		/* no-op */
	}

	/**
	 * Post-processes a test after execution of after
	 * lifecycle callbacks of the underlying test framework — for example,
	 * by tearing down test fixtures.
	 * 
This method must be called immediately after
	 * framework-specific after lifecycle callbacks. For historical
	 * reasons, this method is named {@code afterTestMethod}. Since the
	 * introduction of {@link #afterTestExecution}, a more suitable name for
	 * this method might be something like {@code afterTestTearDown} or
	 * {@code afterEach}; however, it is unfortunately impossible to rename
	 * this method due to backward compatibility concerns.
	 * 
The default implementation is empty. Can be overridden by
	 * concrete classes as necessary.
	 * @param testContext the test context in which the test method was
	 * executed; never {@code null}
	 * @throws Exception allows any exception to propagate
	 * @see #beforeTestMethod
	 * @see #beforeTestExecution
	 * @see #afterTestExecution
	 */
	default void afterTestMethod(TestContext testContext) throws Exception {
		/* no-op */
	}

	/**
	 * Post-processes a test class after execution of all tests within
	 * the class.
	 * 
This method should be called immediately after framework-specific
	 * after class lifecycle callbacks.
	 * 
The default implementation is empty. Can be overridden by
	 * concrete classes as necessary.
	 * @param testContext the test context for the test; never {@code null}
	 * @throws Exception allows any exception to propagate
	 * @since 3.0
	 */
	default void afterTestClass(TestContext testContext) throws Exception {
		/* no-op */
	}

由于本文今天关注的重点是如何实现回滚的，那么咱们需要care的其实就是TransactionalTestExecutionListener的beforeTestMethod和afterTestMethod。

3.3 TransactionalTestExecutionListener

3.3.1 beforeTestMethod

public void beforeTestMethod(final TestContext testContext) throws Exception {
		Method testMethod = testContext.getTestMethod();
		Class<?> testClass = testContext.getTestClass();
		Assert.notNull(testMethod, "Test method of supplied TestContext must not be null");

		TransactionContext txContext = TransactionContextHolder.removeCurrentTransactionContext();
		Assert.state(txContext == null, "Cannot start new transaction without ending existing transaction");

		PlatformTransactionManager tm = null;
    // 获取方法或者类上的@Transactional注解信息
		TransactionAttribute transactionAttribute = this.attributeSource.getTransactionAttribute(testMethod, testClass);
		// 如果加了@Transactional注解
		if (transactionAttribute != null) {
			transactionAttribute = TestContextTransactionUtils.createDelegatingTransactionAttribute(testContext,
				transactionAttribute);

			if (transactionAttribute.getPropagationBehavior() == TransactionDefinition.PROPAGATION_NOT_SUPPORTED) {
				return;
			}
// transactionManager默认情况下会获取容器里的PlatformTransactionManager
			tm = getTransactionManager(testContext, transactionAttribute.getQualifier());
			Assert.state(tm != null,
					() -> "Failed to retrieve PlatformTransactionManager for @Transactional test: " + testContext);
		}
		if (tm != null) {
      // 需要重点关注isRollback方法，新建了一个事务上下文txContext
			txContext = new TransactionContext(testContext, tm, transactionAttribute, isRollback(testContext));
			runBeforeTransactionMethods(testContext);
      // 这里面通过TransactionManager的getTransaction获取了对应的transactionStatus，然后设置到了transactionContext上了
			txContext.startTransaction();
      // 将transactionContext放到了ThreadLocal上
			TransactionContextHolder.setCurrentTransactionContext(txContext);
		}
	}

3.3.3.1 isRollback

protected final boolean isRollback(TestContext testContext) throws Exception {
  // 如果类上面没有加@Rollback注解，那么此处rollback = true；只有@Rollback(value=false)，此处才是false
		boolean rollback = isDefaultRollback(testContext);
  // 获取方法上的@Rollback注解
		Rollback rollbackAnnotation =
				AnnotatedElementUtils.findMergedAnnotation(testContext.getTestMethod(), Rollback.class);
		if (rollbackAnnotation != null) {
      // 返回方法上的@Rollback注解的value
			boolean rollbackOverride = rollbackAnnotation.value();
			rollback = rollbackOverride;
		}
		else {
		}
  // 所以此处默认是true
		return rollback;
	}

3.3.1.2 isDefaultRollback

protected final boolean isDefaultRollback(TestContext testContext) throws Exception {
		Class<?> testClass = testContext.getTestClass();
		Rollback rollback = AnnotatedElementUtils.findMergedAnnotation(testClass, Rollback.class);
		boolean rollbackPresent = (rollback != null);
		if (rollbackPresent) {
			boolean defaultRollback = rollback.value();
			return defaultRollback;
		}
		// else
		return true;
	}

3.3.3.3 TransactionContext

TransactionContext(TestContext testContext, PlatformTransactionManager transactionManager,
			TransactionDefinition transactionDefinition, boolean defaultRollback) {
		this.testContext = testContext;
		this.transactionManager = transactionManager;
		this.transactionDefinition = transactionDefinition;
		this.defaultRollback = defaultRollback;
  // 因此此处的flaggedForRollback被设置成了true
		this.flaggedForRollback = defaultRollback;
	}

综上，可以发现，在beforeTestMethod中，会读取方法的@Transactional注解，如果类或者方法上没有@Rollback，那么默认的会设置一个 TransactionContext，并将其flaggedForRollback设置为true

3.3.2 afterTestMethod

public void afterTestMethod(TestContext testContext) throws Exception {
		Method testMethod = testContext.getTestMethod();
		Assert.notNull(testMethod, "The test method of the supplied TestContext must not be null");
		// 从ThreadLocal上拿到了事务上下文txContext
		TransactionContext txContext = TransactionContextHolder.removeCurrentTransactionContext();
		// If there was (or perhaps still is) a transaction...
		if (txContext != null) {
			TransactionStatus transactionStatus = txContext.getTransactionStatus();
			try {
				// If the transaction is still active...
        // 获取事务状态transactionStatus，
				if (transactionStatus != null && !transactionStatus.isCompleted()) {
          // 在该方法里进行回滚
					txContext.endTransaction();
				}
			}
			finally {
				runAfterTransactionMethods(testContext);
			}
		}
	}

3.3.2.1 endTransaction

void endTransaction() {
		Assert.state(this.transactionStatus != null,
				() -> "Failed to end transaction - transaction does not exist: " + this.testContext);

		try {
			if (this.flaggedForRollback) {
        // 进入此处，调用了transactionManager的rollback方法
				this.transactionManager.rollback(this.transactionStatus);
			}
			else {
				this.transactionManager.commit(this.transactionStatus);
			}
		}
		finally {
			this.transactionStatus = null;
		}
	}

3.4 小结

大致流程如下

1. 执行beforeTestMethod，封装了一个TransactionContext，并把其中的flaggedForRollback设置成true，并把TransactionContext放到ThreadLocal里
2. 然后执行我们的单测方法
3. 执行afterTestMethod，从ThreadLocal里取出TransactionContext，判断里面的flaggedForRollback，如果为true，就rollback，否则就commit

在txContext.startTransaction();中，会调用transactionManager.getTransaction来获取当前的事务内容transactionStatus，关于这块的内容，可以参见这篇博客

总结

Spring这个盒子，很多实现方式都是类似的，特别这个SPI机制，很多地方使用，也是Spring框架开放给第三方框架用来接入的一个大杀器。