스프링 통합 테스트를 위한 AbstractTestExecutionListener

개요

테스트는 메소드별로 타 테스트에 영향이 없어야한다. 스프링에서는 @Transactional 어노테이션을 사용하면 테스트 실행 시 발생한 DB작업에 대해 테스트가 종료 시 자동으로 롤백해주는 기능을 제공한다. 지금까지 해당 기능을 통해 통합 테스트를 수행했고 별다른 문제가 발생하지 않았는데, 이번에 비즈니스 로직을 추가하면서 전파 속성이 REQUIRES_NEW인 트랜잭션 단위에 대해 문제가 발생했고 이에 따라 데이터베이스를 초기화시켜주는 작업을 @Transactional 사용에서 AbstractTestExecutionListener를 사용하는 것으로 변경하게 되었다. 본 포스팅에서는 MySQL 대상으로 각 테스트 시 데이터베이스 초기화를 통해 테스트를 격리시키는 방법에 대해 기술한다.

AbstractTestExecutionListener

AbstractTestExecutionListener를 상속받은 클래스를 구현하고, 내부 메소드를 오버라이드하면 JUnit의 @BeforeEach, @AfterEach와 같은 훅을 걸어줄 수 있다. 

public class IntegrationTestExecutionListener extends AbstractTestExecutionListener {
    private JdbcTemplate getJdbcTemplate(TestContext testContext) {
        return testContext.getApplicationContext().getBean(JdbcTemplate.class);
    }
}

먼저 AbstractTestExecutionListener를 상속받고, 위처럼 JdbcTemplate을 얻어오는 메소드를 하나 생성한다. 리스너에는 기본적으로 TestContext 객체가 인자로 들어오기 때문에 스프링 컨테이너에 접근하여 일반적으로 빈을 가져올 때 사용하는 getBean() 메소드도 사용할 수 있다. 따라서 해당 명령어를 통해 현재 스프링 컨테이너에 등록되어있는 JdbcTemplate을 가져온다.

private static final Set<String> excludedTableNames =
        new HashSet<>(
                List.of(
                        "flyway_schema_history",
                        "BATCH_JOB_EXECUTION",
                        "BATCH_JOB_EXECUTION_CONTEXT",
                        "BATCH_JOB_EXECUTION_PARAMS",
                        "BATCH_JOB_EXECUTION_SEQ",
                        "BATCH_JOB_INSTANCE",
                        "BATCH_JOB_SEQ",
                        "BATCH_STEP_EXECUTION",
                        "BATCH_STEP_EXECUTION_CONTEXT",
                        "BATCH_STEP_EXECUTION_SEQ"));

다음으로 Truncate 제외 대상 테이블을 클래스 내부 필드로 정의한다. 위에서는 flyway, spring batch 관련 테이블을 넣어줬다.

private List<String> getAllTableNames(JdbcTemplate jdbcTemplate) throws SQLException {
    List<String> tableNames = new ArrayList<>();

    try (Connection connection = jdbcTemplate.getDataSource().getConnection()) {
        DatabaseMetaData metaData = connection.getMetaData();
        ResultSet tables = metaData.getTables(null, null, "%", new String[] {"TABLE"});
        while (tables.next()) {
            String tableName = tables.getString("TABLE_NAME");
            if (excludedTableNames.contains(tableName)) {
                continue;
            }
            tableNames.add(tableName);
        }
    } catch (NullPointerException e) {
        log.warn("IntegrationTestExecutionListener | skip truncate. ex={}", e.getMessage());
    }
    return tableNames;
}

다음으로 Truncate시킬 모든 테이블의 이름을 조회한다. catch문에서 NullPointerException을 잡아서 처리하는 부분은, 타 테스트에서 @SpringBootTest를 사용하고 JdbcTemplate를 모킹하는 테스트가 존재하기 때문에 추가한 부분이다. 

참고로 위처럼 try ~ resource 구문을 활용하지 않는 경우 반드시 jdbcTemplate을 통해 얻어온 connection을 close() 해줘야 한다. 그렇지 않으면 테스트를 수행하다가 커넥션 부족 문제로 인해 커넥션 타임아웃 문제가 발생한다. 

public interface Connection  extends Wrapper, AutoCloseable { }

(커넥션 클래스의 경우 위처럼 AutoCloseable 클래스를 상속받고 있기 때문에 try ~ resource 구문에 활용할 수 있다)

private void truncateAllTables(
        @NonNull final List<String> tableNames, @NonNull final JdbcTemplate jdbcTemplate) {
    tableNames.forEach(
            table -> {
                jdbcTemplate.execute("TRUNCATE TABLE " + table);
            });
}

다음으로 가져온 테이블을 이름 기반으로 모두 Truncate해주는 메소드를 작성한다.

@Override
public void afterTestMethod(TestContext testContext) throws Exception {
    JdbcTemplate jdbcTemplate = getJdbcTemplate(testContext);
    List<String> tableNames = getAllTableNames(jdbcTemplate);
    truncateAllTables(tableNames, jdbcTemplate);
}

나는 각 테스트 종료 시 데이터베이스를 초기화시키길 원하기 때문에 리스너에서 제공하는 afterTestMethod() 메소드를 오버라이드 하였다. 순서대로 JdbcTemplate을 얻어오고 모든 테이블 이름을 조회하여 Truncate시키는 작업이다.

@Slf4j
public class IntegrationTestExecutionListener extends AbstractTestExecutionListener {
    private static final Set<String> excludedTableNames =
            new HashSet<>(
                    List.of(
                            "flyway_schema_history",
                            "BATCH_JOB_EXECUTION",
                            "BATCH_JOB_EXECUTION_CONTEXT",
                            "BATCH_JOB_EXECUTION_PARAMS",
                            "BATCH_JOB_EXECUTION_SEQ",
                            "BATCH_JOB_INSTANCE",
                            "BATCH_JOB_SEQ",
                            "BATCH_STEP_EXECUTION",
                            "BATCH_STEP_EXECUTION_CONTEXT",
                            "BATCH_STEP_EXECUTION_SEQ"));

    @Override
    public void afterTestMethod(TestContext testContext) throws Exception {
        JdbcTemplate jdbcTemplate = getJdbcTemplate(testContext);
        List<String> tableNames = getAllTableNames(jdbcTemplate);
        truncateAllTables(tableNames, jdbcTemplate);
    }

    private void truncateAllTables(
            @NonNull final List<String> tableNames, @NonNull final JdbcTemplate jdbcTemplate) {
        tableNames.forEach(
                table -> {
                    jdbcTemplate.execute("TRUNCATE TABLE " + table);
                });
    }

    private JdbcTemplate getJdbcTemplate(TestContext testContext) {
        return testContext.getApplicationContext().getBean(JdbcTemplate.class);
    }

    private List<String> getAllTableNames(JdbcTemplate jdbcTemplate) throws SQLException {
        List<String> tableNames = new ArrayList<>();

        try (Connection connection = jdbcTemplate.getDataSource().getConnection()) {
            DatabaseMetaData metaData = connection.getMetaData();
            ResultSet tables = metaData.getTables(null, null, "%", new String[] {"TABLE"});
            while (tables.next()) {
                String tableName = tables.getString("TABLE_NAME");
                if (excludedTableNames.contains(tableName)) {
                    continue;
                }
                tableNames.add(tableName);
            }
        } catch (NullPointerException e) {
            log.warn("IntegrationTestExecutionListener | skip truncate. ex={}", e.getMessage());
        }
        return tableNames;
    }
}

리스너쪽 전체 소스는 다음과 같다.

TestExecutionListeners

@TestExecutionListeners(
        value = {IntegrationTestExecutionListener.class},
        mergeMode = TestExecutionListeners.MergeMode.MERGE_WITH_DEFAULTS)

이제 위에서 생성한 리스너를 위처럼 등록시켜주면 각 테스트가 종료될 때 마다 모든 테이블을 Truncate 시키는 모습을 확인할 수 있다.