W tym artykule poza podaniem na tacy dwóch rozwiązań tytułowego problemu, porozmawiamy sobie również o kilku tematach pobocznych. Uważam je za warte omówienia, ponieważ wyjaśniają specyficzne zachowanie SpringBoot i bazy H2 podczas uruchamiania testów JUnit.

TL;DR

Gotowe rozwiązanie znajdziesz tu oraz na moim GitLabie. Zachęcam jednak do przeczytania całości, ponieważ dochodząc do rozwiązania zahaczam o wiele ciekawych tematów z zakresu SpringBoot oraz samej bazy danych H2.

Dlaczego w ogóle czyścić bazę danych podczas testów?

Resetowanie stanu bazy danych przed uruchomieniem każdego testu może być szczególnie przydatne, gdy mamy do czynienia z testami integracyjnymi oraz Springiem. Bardzo często w przypadku testów integracyjnych wykorzystujemy bazę danych H2 w trybie in-memory, co pozwala nam na testowanie warstwy JPA z prawdziwą bazą danych. Problem w tym, że owa instancja bazy powoływana jest wraz z podnoszeniem kontekstu Springa i przechowuje ona swój stan aż do wyłączenia aplikacji. Powoduje to problemy z ustaleniem stanu w jakim znajduje się baza przed rozpoczęciem danego testu i czy pierwotny stan nie wpływa wynik testu. Dlatego też dużo wygodniej byłoby wyczyścić bazę danych przed uruchomieniem każdego przypadku testowego, tak aby mieć pewność, że za każdym razem wykonuje się on w tych samych warunkach. W dalszej części przedstawię dwie metody resetowania stanu bazy oraz omówię wady i zalety każdej z nich.

Przypadek testowy

Przygotowałem poniższy test, aby lepiej zobrazować opisywany problem. Zauważ, że mamy w nim dwie metody testowe. Każda z nich dodaje książkę do repozytorium książek oraz sprawdza czy repozytorium zawiera dokładnie jedną książkę.

package com.bettercoding.h2databasecleanup;

import com.bettercoding.h2databasecleanup.db.Book;
import com.bettercoding.h2databasecleanup.helper.BookRepositoryHelper;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;

import java.util.UUID;

import static org.junit.jupiter.api.Assertions.assertEquals;

@SpringBootTest
class H2DatabaseCleanupTests {

    @Autowired
    private BookRepositoryHelper bookRepositoryHelper;


    @Test
    void firstTest() {
        //given
        Book book = new Book(UUID.randomUUID().toString(), "A Passage to India", "E.M. Foster");
        //when
        bookRepositoryHelper.addBook(book);
        // then
        assertEquals(1, bookRepositoryHelper.findAllBooks().size());
    }

    @Test
    void secondTest() {
        //given
        Book book = new Book(UUID.randomUUID().toString(), "A Revenue Stamp", "Amrita Pritam");
        //when
        bookRepositoryHelper.addBook(book);
        // then
        assertEquals(1, bookRepositoryHelper.findAllBooks().size());
    }


}

package com.bettercoding.h2databasecleanup.helper;

import com.bettercoding.h2databasecleanup.db.Book;
import com.bettercoding.h2databasecleanup.db.BookRepository;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.transaction.Transactional;
import java.util.List;

@Component
public class BookRepositoryHelper {
    @Autowired
    BookRepository bookRepository;

    @Transactional
    public void addBook(Book book) {
        bookRepository.save(book);
    }

    @Transactional
    public List<Book> findAllBooks() {
        return bookRepository.findAll();
    }
}

package com.bettercoding.h2databasecleanup.db;

import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.NoArgsConstructor;
import lombok.Setter;

import javax.annotation.processing.Generated;
import javax.persistence.Column;
import javax.persistence.Entity;
import javax.persistence.Id;

@Getter
@Setter
@NoArgsConstructor
@Entity
@AllArgsConstructor
public class Book {
    @Id
    private String id;

    @Column
    private String title;

    @Column
    private String author;
}

package com.bettercoding.h2databasecleanup.db;

import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.stereotype.Repository;

@Repository
public interface BookRepository extends JpaRepository<Book, String> {
}

Po uruchomieniu testu okazuje się, że drugi przypadek testowy kończy się niepowodzeniem, ponieważ w repozytorium oprócz nowo dodawanej książki znajduje się już inna – dodana w pierwszym przypadku testowym. Oczywiście możemy napisać ten test w taki sposób aby był odporny na już istniejące dane, jednak w przypadku bardziej złożonych testów wygodniej jest startować z czystą aplikacją i zmodelować sobie dokładnie takie warunki jakie chcemy.

Metoda 1: @DirtiestContext

SpringBoot posiada mechanizm re-używania kontekstów na potrzeby testów. Nie będę się tu zagłębiać nad jego dokładnym opisem ponieważ nie jest to w zakresie tego artykułu. Ważne aby wiedzieć, że jeśli jakaś klasa testowa nie różni się znacznie od innej (np. nie posiada dodatkowych beanów, lub innych elementów konfiguracyjnych), to SpringBoot potrafi bardzo szybko powołać czystą kopię kontekstu ze specjalnego cache’a zamiast startować całą aplikację na nowo. W naszym przykładzie mamy do czynienia dokładnie z taką sytuacją, więc SpringBoot przy uruchamianiu pierwszego przypadku testowego wystartuje kontekst i doda go do cache’a, a następnie wykona pierwszy przypadek. Gdy uruchomi się drugi przypadek testowy, to SpringBoot dostarczy nam czystą kopię kontekstu z cache’a – co jest bardzo szybkie – zamiast startować całą aplikacje na nowo. Niestety baza H2 nie jest powoływana na nowo.

Możemy wymusić na SpringBoot aby ten nie korzystał ze wspomnianego mechanizmu cache’owania kontekstu poprzez dodanie adnotacji @DirtiestContext do klasy testowej. Ważne jest również, aby wymusić startowanie nowego kontekstu dla każdego przypadku testowego poprzez dodanie @DirtiesContext(classMode = DirtiesContext.ClassMode.BEFORE_EACH_TEST_METHOD).

package com.bettercoding.h2databasecleanup;

import com.bettercoding.h2databasecleanup.db.Book;
import com.bettercoding.h2databasecleanup.helper.BookRepositoryHelper;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.annotation.DirtiesContext;

import java.util.UUID;

import static org.junit.jupiter.api.Assertions.assertEquals;

@SpringBootTest
@DirtiesContext(classMode = DirtiesContext.ClassMode.BEFORE_EACH_TEST_METHOD)
class H2DatabaseCleanupTests1 {

    @Autowired
    private BookRepositoryHelper bookRepositoryHelper;


    @Test
    void firstTest() {
        //given
        Book book = new Book(UUID.randomUUID().toString(), "A Passage to India", "E.M. Foster");
        //when
        bookRepositoryHelper.addBook(book);
        // then
        assertEquals(1, bookRepositoryHelper.findAllBooks().size());
    }

    @Test
    void secondTest() {
        //given
        Book book = new Book(UUID.randomUUID().toString(), "A Revenue Stamp", "Amrita Pritam");
        //when
        bookRepositoryHelper.addBook(book);
        // then
        assertEquals(1, bookRepositoryHelper.findAllBooks().size());
    }


}

To rozwiązanie jest bardzo proste w implementacji. Niestety dla każdego przypadku testowego aplikacja uruchamiana jest ponownie, co w zależności od skomplikowania aplikacji może trwać od kilku do kilkudziesięciu sekund. Może to skutkować tym, że gdy liczba przypadków znacznie wzrośnie, to na wynik testów będziemy czekać długie minuty, a nawet godziny. Zwróć proszę uwagę na czasy wykonania obu przypadków testowych, które porównamy z drugą omawianą metodą.

Metoda 2: TestExecutionListener czyszczący bazę danych

Zastanówmy się nad innym sposobem resetowania stanu bazy danych przed uruchomieniem każdego przypadku testowego, tak aby nie rezygnować z zalet cache’owania kontekstu. Tu z pomocą przychodzą nam interfejs TestExecutionListener oraz adnotacja @TestExecutionListeners, które pozwala wpiąć się w cykl wykonywania testów JUnit oraz uruchomienie własnego kodu np. przed uruchomieniem przypadku testowego.

W dalszej części zobaczysz kod źródłowy dwóch klas CleanupH2DatabaseTestListener oraz CleanupH2DbService. Pierwsza z nich odpowiedzialna jest za wpięcie się w cykl wykonania testów w taki sposób, aby zawołać kod odpowiedzialny za czyszczenie bazy danych, który został zaimplementowany w komponencie CleanupH2DbService. Omawiane rozwiązanie można zaprezentować w poniższy sposób:

W procesie usuwania danych w relacyjnej bazie danych pojawia się pewien problem. Mianowicie, nie możemy usunąć wiersza nadrzędnego dopóki istnieją wiersze odwołujące się do niego. W praktyce oznacza to czyszczenie tabel w odpowiedniej kolejności (w przypadku klucza obcego do tej samej tabeli istnieje nawet konieczność usuwania wierszy w odpowiedniej kolejności).

Na szczęście, możemy wykorzystać pewną sztuczkę i wyłączyć na moment sprawdzanie Constraint’ów, dzięki czemu, będziemy stanie usunąć dane w dowolnej kolejności. Po zakończeniu usuwania włączamy z powrotem sprawdzanie Constraint’ów. Dodatkowym krokiem, który warto wykonać jest resetowanie sekwencji.

Ostatnią kwestą, która pozostaje nam do zrobienia to dodanie nowo-napisanego CleanupH2DatabaseTestListener do naszego testu za pomocą adnotacji @TestExecutionListeners. Tu również czai się mały kruczek – aby nadal działało DependencyInjection, to do listy customowych listenerów musimy dodać jeszcze DependencyInjectionTestExecutionListener. Zobaczmy zatem finalne rozwiązanie, które może również znaleźć na moim GitLabie.

package com.bettercoding.h2databasecleanup.listener;

import lombok.extern.slf4j.Slf4j;
import org.springframework.core.Ordered;
import org.springframework.test.context.TestContext;
import org.springframework.test.context.TestExecutionListener;

@Slf4j
public class CleanupH2DatabaseTestListener implements TestExecutionListener, Ordered {

    private static final String H2_SCHEMA_NAME = "PUBLIC";

    @Override
    public void beforeTestMethod(TestContext testContext) throws Exception {
        TestExecutionListener.super.beforeTestMethod(testContext);
        cleanupDatabase(testContext);
    }

    private void cleanupDatabase(TestContext testContext) {
        log.info("Cleaning up database begin");
        CleanupH2DbService cleanupH2DbService = testContext.getApplicationContext().getBean(CleanupH2DbService.class);
        cleanupH2DbService.cleanup(H2_SCHEMA_NAME);
        log.info("Cleaning up database end");
    }

    @Override
    public int getOrder() {
        return 0;
    }
}

package com.bettercoding.h2databasecleanup.listener;

import lombok.RequiredArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import javax.transaction.Transactional;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.Statement;
import java.util.HashSet;
import java.util.Set;

@Component
@Slf4j
@RequiredArgsConstructor
public class CleanupH2DbService {
    public static final String H2_DB_PRODUCT_NAME = "H2";
    private final DataSource dataSource;

    @SneakyThrows
    @Transactional(Transactional.TxType.REQUIRES_NEW)
    public void cleanup(String schemaName) {
        try (Connection connection = dataSource.getConnection();
             Statement statement = connection.createStatement()) {
            if (isH2Database(connection)) {
                disableConstraints(statement);
                truncateTables(statement, schemaName);
                resetSequences(statement, schemaName);
                enableConstraints(statement);
            } else {
                log.warn("Skipping cleaning up database, because it's not H2 database");
            }
        }
    }

    private void resetSequences(Statement statement, String schemaName) {
        getSchemaSequences(statement, schemaName).forEach(sequenceName ->
                executeStatement(statement, String.format("ALTER SEQUENCE %s RESTART WITH 1", sequenceName)));
    }

    private void truncateTables(Statement statement, String schemaName) {
        getSchemaTables(statement, schemaName)
                .forEach(tableName -> executeStatement(statement, "TRUNCATE TABLE " + tableName));
    }

    private void enableConstraints(Statement statement) {
        executeStatement(statement, "SET REFERENTIAL_INTEGRITY TRUE");
    }

    private void disableConstraints(Statement statement) {
        executeStatement(statement, "SET REFERENTIAL_INTEGRITY FALSE");
    }

    @SneakyThrows
    private boolean isH2Database(Connection connection) {
        return H2_DB_PRODUCT_NAME.equals(connection.getMetaData().getDatabaseProductName());
    }

    @SneakyThrows
    private void executeStatement(Statement statement, String sql) {
        statement.executeUpdate(sql);
    }

    @SneakyThrows
    private Set<String> getSchemaTables(Statement statement, String schemaName) {
        String sql = String.format("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES  where TABLE_SCHEMA='%s'", schemaName);
        return queryForList(statement, sql);
    }

    @SneakyThrows
    private Set<String> getSchemaSequences(Statement statement, String schemaName) {
        String sql = String.format("SELECT SEQUENCE_NAME FROM INFORMATION_SCHEMA.SEQUENCES WHERE SEQUENCE_SCHEMA='%s'", schemaName);
        return queryForList(statement, sql);
    }

    @SneakyThrows
    private Set<String> queryForList(Statement statement, String sql) {
        Set<String> tables = new HashSet<>();
        try (ResultSet rs = statement.executeQuery(sql)) {
            while (rs.next()) {
                tables.add(rs.getString(1));
            }
        }
        return tables;
    }
}

package com.bettercoding.h2databasecleanup;

import com.bettercoding.h2databasecleanup.db.Book;
import com.bettercoding.h2databasecleanup.helper.BookRepositoryHelper;
import com.bettercoding.h2databasecleanup.listener.CleanupH2DatabaseTestListener;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.test.context.TestExecutionListeners;
import org.springframework.test.context.support.DependencyInjectionTestExecutionListener;

import java.util.UUID;

import static org.junit.jupiter.api.Assertions.assertEquals;

@SpringBootTest
@TestExecutionListeners(listeners = {DependencyInjectionTestExecutionListener.class, CleanupH2DatabaseTestListener.class})
class H2DatabaseCleanupTests2 {

    @Autowired
    private BookRepositoryHelper bookRepositoryHelper;


    @Test
    void firstTest() {
        //given
        Book book = new Book(UUID.randomUUID().toString(), "A Passage to India", "E.M. Foster");
        //when
        bookRepositoryHelper.addBook(book);
        // then
        assertEquals(1, bookRepositoryHelper.findAllBooks().size());
    }

    @Test
    void secondTest() {
        //given
        Book book = new Book(UUID.randomUUID().toString(), "A Revenue Stamp", "Amrita Pritam");
        //when
        bookRepositoryHelper.addBook(book);
        // then
        assertEquals(1, bookRepositoryHelper.findAllBooks().size());
    }


}

Pozostaje jeszcze tylko uruchomić testy i sprawdzić czy wynik pokrywa się z oczekiwanym. Zwróć proszę uwagę o ile szybciej wykonały się testy w porównaniu z poprzednią metodą. Czyszczenie bazy danych trwało w tym przypadku poniżej 4ms, co jest świetnym wynikiem w porównaniu z poprzednią metodą, która zajęła około 130ms. Należy jednak pamiętać, że pierwsza metoda, to narzut związany ze startem aplikacji i jest zależy od liczby i czasu ładowania beanów, a druga metoda zależy jedynie od liczby tabel i ilości danych. Chcę przez to powiedzieć, że wraz ze wzrostem skomplikowania aplikacji pierwsza metoda znacznie spowolni, podczas gdy druga – niezależnie od warunków – pozostanie tak samo szybka.

To już wszystko co dla Ciebie przygotowałem w tym poradniku. Udostępnij proszę ten post aby pomóc mi dotrzeć do większego grona odbiorców. Dzięki i do zobaczenia w kolejnym poradniku.