Resetting Database Between Spring Integration Tests

When tasked with having to write an integration test or a Spring Webflux test that uses a database, it can be cumbersome to have to reset the database between each test by using @DirtiesContext. Using @DirtiesContext slows down your build significantly since everything, including beans, needs to be reloaded between tests. This post will explain how you can reset the database between tests. I created a sample Spring Boot project if you’d like to check out a working simple example.

The main issue of trying to reset the database is to find out what tables exist without hardcoding anything. This can be accomplished by looking at the Autowired EntityManager and then running truncate commands. Let’s start with creating a service that will do what we want called TestDbService. Since this is a dangerous service, we’re going to place this in the test module and not the actual code module.

@RequiredArgsConstructor
@Service
public class TestDbService {

    // Autowired via constructor based injected (using Lombok for generated constructor)
    private final EntityManager entityManager;
    private List<String> tableNames;

    @PostConstruct
    void afterPropertiesSet() {
        tableNames = entityManager.getMetamodel().getEntities().stream()
            .filter(entityType -> entityType.getJavaType().getAnnotation(Table.class) != null)
            .map(entityType -> entityType.getJavaType().getAnnotation(Table.class))
            .map(this::convertToTableName) // TODO
            .collect(Collectors.toUnmodifiableList());
    }
}

First we’re going to query the metadata and grab everything from the @Table annotations. Then we’re going to convert the name that we need to use for h2 database. Depending on how you make your tables, you may not need this but If you have a table named CamelCaseTable, then h2 is going to create a table named Camel_Case_Table, so we can’t use the name you use as-is for the truncate commands. We need our truncate statement to reflect the correct name. We can fix this with some regex in the convertToTableName(Table table) function below:

    /**
     * Converts an (optional) schema and table on a {@link Table} annotation to something that h2
     * uses when it generates tables.
     */
    private String convertToTableName(Table table) {
        String schema = table.schema();
        String tableName = table.name();

        String convertedSchema = StringUtils.hasText(schema) ? schema.toLowerCase() + "." : "";
        String convertedTableName = tableName.replaceAll("([a-z])([A-Z])", "$1_$2");

        return convertedSchema + convertedTableName;
    }

Using the regex make it now we know tables that the h2 database in actually using. The schema part of optional in case you define the schema in your @Table annotations. This will convert it to Schema.Table_name.

Now that we have the tables to convert, we just need to delete the entries. We’ll do this with a new service method named resetDatabase().

    @Transactional
    public void resetDatabase() {
        entityManager.createNativeQuery("SET REFERENTIAL_INTEGRITY FALSE").executeUpdate();

        for (String tableName : tableNames) {
            entityManager.createNativeQuery("TRUNCATE TABLE " + tableName).executeUpdate();
        }

        entityManager.createNativeQuery("SET REFERENTIAL_INTEGRITY TRUE").executeUpdate();
    }

Here we’re getting away with truncating the tables in any order with the referential integrity disablement and reenablement.

Now to bring it all together, we need a test that calls this new service. This is accomplished by calling the service in an @AfterEach method. The example test below is just running the same code twice with no dirty context calls. Without calling the TestDbService, one of the tests would fail since the number of users found would be 4 instead of 2.

/**
 * This test class would be much better served as a {@link DataJpaTest}. (where the TestDbService wouldn't be needed). This is a {@link
 * SpringBootTest} only to show a simple test of using the {@link TestDbService}.
 */
@SpringBootTest
class ImportantUserRepositoryTest {

    @Autowired
    private ImportantUserRepository importantUserRepository;

    @Autowired
    private TestDbService testDbService;

    private ImportantUser user1;
    private ImportantUser user10;
    private ImportantUser user20;

    @BeforeEach
    void beforeEach() {
        user1 = ImportantUser.builder()
            .firstName("first10")
            .lastName("last")
            .build();

        user10 = ImportantUser.builder()
            .firstName("first")
            .lastName("last10")
            .build();

        user20 = ImportantUser.builder()
            .firstName("first20")
            .lastName("last20")
            .build();

        user1 = importantUserRepository.save(user1);
        user10 = importantUserRepository.save(user10);
        user20 = importantUserRepository.save(user20);
    }

    @AfterEach
    void afterEach() {
        testDbService.resetDatabase();
    }

    @Test
    void findByFistNameLikeOrLastNameLike() {
        List<ImportantUser> foundUsers = importantUserRepository.findByFistNameLikeOrLastNameLike("10");

        assertThat(foundUsers, hasSize(2));
        assertThat(foundUsers.get(0).getId(), is(user1.getId()));
        assertThat(foundUsers.get(1).getId(), is(user10.getId()));
    }

    @Test
    void findByFistNameLikeOrLastNameLike2() {
        List<ImportantUser> foundUsers = importantUserRepository.findByFistNameLikeOrLastNameLike("10");

        assertThat(foundUsers, hasSize(2));
        assertThat(foundUsers.get(0).getId(), is(user1.getId()));
        assertThat(foundUsers.get(1).getId(), is(user10.getId()));
    }
}

And that’s it! Again check out the Github sample project if you’d like to see a working example.

About the Author

Jeff Torson profile.

Jeff Torson

Principal Technologist

Jeff is a full stack developer with experience in the government/defense and finance industry. He has experience ranging from thick client Eclipse RCP programs to microservices using Spring Boot for data access and Elasticsearch. He enjoys learning about anything related to the IT field and has even managed Linux and Windows servers and setup deployment pipelines. He is a true believer in that if something is worth doing, then it’s worth doing right the first time and fully unit/integration tested.

Leave a Reply

Your email address will not be published. Required fields are marked *

Related Blog Posts
An Exploration in Rust: Musings From a Java/C++ Developer
Why Rust? It’s fast (runtime performance) It’s small (binary size) It’s safe (no memory leaks) It’s modern (build system, language features, etc) When Is It Worth It? Embedded systems (where it is implied that interpreted […]
Getting Started with CSS Container Queries
For as long as I’ve been working full-time on the front-end, I’ve heard about the promise of container queries and their potential to solve the majority of our responsive web design needs. And, for as […]
Simple improvements to making decisions in teams
Software development teams need to make a lot of decisions. Functional requirements, non-functional requirements, user experience, API contracts, tech stack, architecture, database schemas, cloud providers, deployment strategy, test strategy, security, and the list goes on. […]
JavaScript Bundle Optimization – Polyfills
If you are lucky enough to only support a small subset of browsers (for example, you are targeting a controlled set of users), feel free to move along. However, if your website is open to […]