Spring Service 계층 테스트와 데이터 초기화

JPA는 특성상 실제 쿼리를 날리지 않아도 DB의 값이 변경되는 경우가 많습니다. 그리고 보통 영속성 컨택스트의 생명 주기가 서비스 계층과 비슷하기 때문에 서비스 계층 단위 테스트를 할 때는 JPA를 포함시켜 테스트를 많이 합니다. 특히 스프링의 @Transactional은 테스트 코드에 추가할 경우, 테스트가 끝나면 자동으로 트랜잭션을 롤백해주기 때문에 간단하게 JPA를 포함시켜 서비스 계층을 테스트할 수 있습니다. 하지만 그랬다가는 예상치 못하는 오류가 발생하는 경우가 있기 때문에 주의해야 합니다.

이번 포스팅에서는 서비스 계층 테스트 코드에 @Transactional을 붙였을 때 문제점과, @Transactional을 이용하지 않고 테스트 코드를 작성하는 방법에 대해 적으려고 합니다.

Spring Boot Transactional 테스트 문제점

Team과 Member가 일대다 양방향 관계에 있다는 가정하에, 멤버의 팀 이름을 가져오는 메서드를 다음과 같이 작성하고 테스트 해보겠습니다. Repository는 Spring Data JPA의 Repository를 사용하였습니다.

@Service
@RequiredArgsConstructor
public class MemberService {

    private final MemberRepository memberRepository;

    @Transactional(readOnly = true)
    public Team getTeamWithLazy(Long memberId) {
        Member member = memberRepository.findById(memberId)
                .orElseThrow(IllegalArgumentException::new);
        return member.getTeam();
    }
}

 

@Transactional
@SpringBootTest
public class TransactionTest {

	...

    @Test
    @DisplayName("멤버의 팀의 이름은 저장된 팀의 이름과 같아야 함")
    void getTeamWithLazyTest() throws Exception {
        //given
        Team team = teamRepository.save(new Team("team1"));
        Member member = memberRepository.save(new Member(team));

        //when
        String result = memberService.getTeamNameWithLazy(member.getId());

        //then
        assertThat(result).isEqualTo(team.getName());
    }
}

 

코드도 언뜻 보기에는 큰 문제 없어보이고 테스트도 무사히 통과했습니다.

그런데 어딘가 조금 이상한 부분이 보입니다. 우선 쿼리를 보면 insert 쿼리만 딱 두 번 나가고 서비스 계층에서 발생했어야 할 쿼리가 하나도 생기지 않았습니다. 왜 이런 결과가 나왔는지 간단하게 설명을 드리겠습니다.

우선 given에서 save를 통해 영속성 컨택스트의 1차 캐시에 저장됩니다. 그 후 when의 서비스 계층으로 들어가 멤버를 조회하고, 지연 로딩을 통해 멤버의 팀과 팀 이름을 조회합니다. 이때, 트랜잭션의 범위가 테스트 코드 전체이기 때문에, given에서 저장한 데이터들이 1차 캐시에 계속 남아있고, 굳이 쿼리를 날리지 않아도 데이터 조회가 가능하게 됩니다.

그렇다면 만약 서비스 계층에 @Transactional을 깜빡하고 빼먹으면 어떻게 될까요? 

@Service
@RequiredArgsConstructor
public class MemberService {

    private final MemberRepository memberRepository;

    //트랜잭션 제거
    //@Transactional(readOnly = true)
    public Team getTeamWithLazy(Long memberId) {
        Member member = memberRepository.findById(memberId)
                .orElseThrow(IllegalArgumentException::new);
        return member.getTeam();
    }
}

Member에서 Team을 지연로딩 설정을 해놨기 때문에 실제 상황이라면 프록시를 초기화할 수 없다는 예외가 발생해야 하지만 테스트 코드는 아까와 똑같은 결과를 내며 테스트를 통과시킵니다.

이유는 이전과 같습니다. 트랜잭션이 데이터를 저장하고 조회하는 부분까지 걸려있기 때문에. 지연 로딩이 발생할 이유 없이 1차 캐시에서 데이터를 가져오면 되기 때문입니다.

즉, 테스트 코드에 단순히 트랜잭션을 추가하면, 프로덕션 상황과 범위가 다른 트랜잭션을 사용하고 있다는 문제가 발생하게 됩니다. 이 외에도 서비스가 복잡해짐에 따라 어떤 값이 db에 저장되고 변경되는지 확인하기 힘들어지는 등 여러 문제가 발생하게 됩니다. 이를 해결하기 위해 given 과정 마지막에 flush를 하는 등의 작업을 추가할 수 있지만, 이는 오히려 테스트 코드의 복잡성과, 잘못된 테스트 코드의 작성 확률을 높이고, 테스트 코드가 서비스 코드와 과하게 결합되어 리팩터링의 어려움을 낳는 결과를 가져옵니다.

그렇기 때문에 테스트 코드에서 트랜잭션을 제거하는 방법을 소개하겠습니다.

@AfterEach를 이용한 데이터 초기화

테스트 코드의 트랜잭션을 제거한다는 뜻은 결국 테스트 후, 데이터를 롤백, 또는 초기화 시킨다는 의미와 같으므로 단순히 테스트에 사용된 엔티티의 모든 데이터를 afterEach구문에서 날려주면 됩니다.

@AfterEach
void clear() {
    memberRepository.deleteAll();
    teamRepository.deleteAll();
}

하지만 서비스가 복잡해지고 엔티티가 여러 다른 엔티티들과 관계를 맺을수록 지워야 하는 엔티티 수도 많아지고, 무엇보다 엔티티 id 값 설정을 IDENTITY 즉, Auto Increment로 했을 경우 Id가 초기화되지 않아 Id를 사용하는 테스트에 문제가 생길 가능성이 높아집니다.

그렇기 때문에 데이터를 초기화 하는 방법으로 SQL을 이용한 방법을 알려드리려고 합니다.

SQL을 이용한 데이터 초기화

SQL을 이용하여 각 테스트 전에 데이터를 모두 초기화하는 방법입니다. 이를 이용하면 각각의 테스트 코드가 트랜잭션에 관해서는 서로 독립되게 만들 수 있습니다. 이를 수행하기 위해 다음과 같은 클래스를 작성합니다.

@Component
public class DatabaseCleaner {

    private static final String TRUNCATE_TABLE = "TRUNCATE TABLE %s";
    private static final String RESET_TABLE_ID = "ALTER TABLE %s ALTER COLUMN ID RESTART WITH 1";
    private static final String REFERENTIAL = "SET REFERENTIAL_INTEGRITY %s";

    @PersistenceContext
    private EntityManager em;

    @Autowired
    private DataSource dataSource;

    private List<String> tableNames = new ArrayList<>();

    @PostConstruct
    public void initTableNames() {
        try (Connection connection = dataSource.getConnection();
             ResultSet tables = connection.getMetaData().getTables(null, "PUBLIC", null, new String[]{"TABLE"})) {
            while (tables.next()) {
                tableNames.add(tables.getString("TABLE_NAME"));
            }
        } catch (Exception e) {
            throw new IllegalArgumentException("테이블 이름 초기화 실패");
        }
    }

    @Transactional
    public void clear() {
        em.clear();

        query(REFERENTIAL, "FALSE");
        for (String tableName : tableNames) {
            query(TRUNCATE_TABLE, tableName);
            query(RESET_TABLE_ID, tableName);
        }
        query(REFERENTIAL, "TRUE");
    }

    private void query(String query, String param) {
        em.createNativeQuery(String.format(query, param)).executeUpdate();
    }
}

이 클래스를 처음 스프링 컨테이너에 등록한 후, 모든 테스트가 끝날 때마다 clear 메서드를 호출하면 데이터가 제거가 됩니다.

해당 코드를 조금 더 자세하게 살펴보겠습니다.

@PostConstruct
public void initTableNames() {
    try (Connection connection = dataSource.getConnection();
         ResultSet tables = connection.getMetaData().getTables(null, "PUBLIC", null, new String[]{"TABLE"})) {
        while (tables.next()) {
            tableNames.add(tables.getString("TABLE_NAME"));
        }
    } catch (Exception e) {
        throw new IllegalArgumentException("테이블 이름 초기화 실패");
    }
}

우선 테이블 이름을 초기화 하는 메서드입니다. 해당 메서드는 @PostConstructor 어노테이션을 붙여, 해당 클래스가 스프링 컨테이너에 등록될 때 이 메서드를 호출하도록 설정합니다.

그리고 dataSource를 이용해 db 커넥션을 가져온 후, 메타 데이터를 이용해 테이블 이름을 가져와 tableNames 리스트에 넣어주는 작업을 해줍니다. 이때, getTables 메서드의 두 번째 인자로 "PUBLIC"을 넣어주는 이유는, 보통 단위 테스트에 사용하는 h2 데이터베이스는 사용자가 직접 만든 테이블을 기본적으로 "PUBLIC"이라는 스키마에 저장하기 때문입니다. 그렇지 않으면 여러 h2 디비 자체 관련 정보 테이블이 함께 넘어오기 때문에, 후에 리스트에 담긴 테이블을 이용해 table truncate를 할 때 오류가 발생합니다.

entity 리스트를 가져와 테이블 이름을 가져오는 방법도 있습니다.

    @PostConstruct
    public void initTableNames() {
        tableNames = em.getMetamodel().getEntities().stream()
                    //디비 테이블 규칙에 맞는 이름
                    //ex .map(entityType -> StringUtils.capitalize(entityType.getName()))
                    .toList();
    }

하지만 이 코드는 @Table(name = "xxx")처럼 따로 테이블 이름을 정하거나, 테이블 이름 규칙이 대문자 또는 스네이크 전략 등 엔티티와 이름이 많이 다를 경우 따로 값을 그에 맞춰줘야 하기 때문에 실제로는 사용하기 힘든 것 같습니다.

이제 데이터를 초기화 하는 코드를 보겠습니다.

private static final String TRUNCATE_TABLE = "TRUNCATE TABLE %s";
private static final String RESET_TABLE_ID = "ALTER TABLE %s ALTER COLUMN ID RESTART WITH 1";
private static final String REFERENTIAL = "SET REFERENTIAL_INTEGRITY %s";

@Transactional
public void clear() {
    em.clear();

    query(REFERENTIAL, "FALSE");
    for (String tableName : tableNames) {
        query(TRUNCATE_TABLE, tableName);
        query(RESET_TABLE_ID, tableName);
    }
    query(REFERENTIAL, "TRUE");
}

private void query(String query, String param) {
    em.createNativeQuery(String.format(query, param)).executeUpdate();
}

이 clear 라는 메서드를 통해 테스트가 끝날 때마다 테이블을 초기화시켜줍니다. 우선 혹시 모를 영속성 캐시 값을 초기화를 합니다. 테이블을 초기화하기 전에 테이블 간의 관계가 있으면 예외가 발생하기 때문에 Referential Integrity(참조 무결성)을 false로 변경을 해준 뒤 테이블을 초기화시킵니다. 이때 테이블의 Id 자동 생성 값도 함께 1로 초기화를 합니다. 그 후 다시 참조 무결성 값을 true로 바꿔주면 테이블 초기화가 끝납니다. 

해당 클래스를 컴포넌트로 등록을 해뒀기 때문에, 서비스 계층 테스트 코드를 작성할 때 해당 컴포넌트를 가져와 @AfterEach 또는 @BeforeEach 에 설정을 해줍니다.

@Autowired
DatabaseCleaner cleaner;

@AfterEach
void clear() {
    cleaner.clear();
}

이를 적용시켜 테스트를 다시 돌려보겠습니다.

드디어 우리가 원하는대로 쿼리가 나가며 테스트를 통과했음을 볼 수 있습니다. 이제는 서비스 계층의 @Transactional 등의 JPA 관련 기능을 프로덕션 상황과 맞춰서 테스트가 가능해졌습니다. 

Junit Extension을 이용한 자동 데이터 초기화

위 코드는 모든 테스트 파일마다 DataCleaner라는 객체를 주입받고, @AfterEach 구문에 clear 메서드를 호출해야 하는 불편함이 있습니다. 이런 번거로움을 Junit의 Extension 기능을 이용해 해결할 수 있습니다. 보통 많은 분들이 단위 테스트를 하면서 Mockito 라이브러리를 사용하기 위해 @ExtendWith(MockitoExtension.class)를 사용해 보셨을 텐데 바로 이때 사용하는 @ExtendWith을 사용할 겁니다.

우선 clear 메서드를 호출할 Extension 클래스를 만들어줍니다. 코드는 아래와 같습니다.

public class DBCleanerExtension implements AfterEachCallback {

    @Override
    public void afterEach(ExtensionContext context) throws Exception {
        DatabaseCleaner databaseCleaner = SpringExtension.getApplicationContext(context)
                .getBean(DatabaseCleaner.class);
        databaseCleaner.clear();
    }
}

내용은 단순합니다. AfterEachCallback을 상속받은 후 afterEach를 오버라이드 해주면, 해당 Extension을 사용한 테스트 코드는 테스트 후 저 afterEach을 호출하게 되고, 결국 우리가 원하는 테스트 후 데이터 초기화 작업을 하게 됩니다.

마지막으로 테스트 코드의 클래스 선언 부분에 해당 Extension을 추가해 주면 됩니다.

@ExtendWith(DBCleanerExtension.class)
@SpringBootTest
public class TransactionTest {
    ...
}

 

이렇게 서비스 계층 테스트를 작성할 때 @Transactional을 그냥 사용하면 생기는 문제점들과 SQL을 이용한 데이터 초기화 방법을 알아봤습니다.

 

참고

https://miensol.pl/clear-database-in-spring-boot-tests/

https://github.com/woowacourse-teams/2022-thankoo/tree/develop/backend/src/test/java/com/woowacourse/thankoo/common/support 

https://velog.io/@tmdgh0221/%EC%8A%A4%ED%94%84%EB%A7%81-%ED%85%8C%EC%8A%A4%ED%8A%B8-%EC%BC%80%EC%9D%B4%EC%8A%A4%EC%97%90%EC%84%9C%EC%9D%98-Transactional-%EC%9C%A0%EC%9D%98%EC%A0%90