Automated database cleanup during integration test run with JUnit5 and Spring Boot

Integration tests are an integral part of software testing, being responsible for joint testing of different parts of the application. Hence, they usually require the application (partially) running with all the attached strings like a database connection. Now, depending on your application, it might be necessary to maintain a clean data sleeve between different integration test classes in oder to reliably and reproducibly test certain behaviours without side effects from residual data. For this, it might come in handy if you simply clean your database e.g. after each test class execution. That is at least my approach to keep my database state known. To achieve this in JUnit5/Spring Boot I do the following:

First I create a singleton component which implements the InitializingBean interface.

@Component
@Profile("integration_test")
public class DatabaseCleanup implements InitializingBean {

    @PersistenceContext
    private EntityManager entityManager;

    private List<String> tableNames;

    @Transactional
    public void execute() {
        this.entityManager.flush();
        this.entityManager.createNativeQuery("SET CONSTRAINTS ALL DEFERRED").executeUpdate();
        for (final String tableName : this.tableNames) {
            this.entityManager.createNativeQuery("ALTER TABLE " + tableName + " DISABLE TRIGGER ALL").executeUpdate();
            this.entityManager.createNativeQuery("TRUNCATE TABLE " + tableName + " CASCADE").executeUpdate();
            this.entityManager.createNativeQuery("ALTER TABLE " + tableName + " ENABLE TRIGGER ALL").executeUpdate();
        }
        this.entityManager.createNativeQuery("SET CONSTRAINTS ALL IMMEDIATE").executeUpdate();
    }

    @Override
    public void afterPropertiesSet() {
        this.tableNames = this.entityManager.getMetamodel().getEntities().stream()
                .filter(e -> e.getJavaType().getAnnotation(Table.class) != null)
                .map(e -> e.getJavaType().getAnnotation(Table.class).name())
                .toList();
    }
}

The InitializingBean interface is necessary because we want our component to wait for its complete initialization including the EntityManager from within the persistence context. Afterwards, the afterPropertiesSet method will populate our class variable tableNames. There are certainly other ways to implement this; One could also simply execute the table name logic within the executemethod. But that does not guarantee in any way that the EntityManager is successfully autowired (except for a proxy). This could be troublesome in case the method is used in any constructor. I personally think this is a precautious and lightweight implementation that does not add any complexity and thus is favored by me.

Now, the logic obtains the names of all tables within my current schema. After that all constraints on the tables are dropped (deferred), which is necessary to just bulk delete stuff. The same needs to be done in regard to any potential triggers on the database to avoid side effects. Finally, after truncating all tables we revert the changes on triggers and constraints and happily move to another test class.

Now to ensure that the database cleanup is executed after each test class we can utilize the @AfterAll annotation from JUnit which we implement either on the test class directly or by inheritance:

@AfterAll
public void cleanupDatabase() {
    this.databaseCleanup.execute();
}

By using this implementation we can easily clean up all the mess we left behind during an integration test on our database. Please be aware that I only used this for Postgres databases. It is possible that you need to adapt the SQL for another database.