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

Object Partners profile.
Leave a Reply

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

Related Blog Posts
Feature Flags in Terraform
Feature flagging any code can be useful to developers but many don’t know how to or even that you can do it in Terraform. Some benefits of Feature Flagging your code You can enable different […]
Infrastructure as Code – The Wrong Way
You are probably familiar with the term “infrastructure as code”. It’s a great concept, and it’s gaining steam in the industry. Unfortunately, just as we had a lot to learn about how to write clean […]
Snowflake CI/CD using Jenkins and Schemachange
CI/CD and Management of Data Warehouses can be a serious challenge. In this blog you will learn how to setup CI/CD for Snowflake using Schemachange, Github, and Jenkins. For access to the code check out […]
How to get your pull requests approved more quickly
TL;DR The fewer reviews necessary, the quicker your PR gets approved. Code reviews serve an essential function on any software codebase. Done right, they help ensure correctness, reliability, and maintainability of code. On many teams, […]