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.

Related Blog Posts
Natively Compiled Java on Google App Engine
Google App Engine is a platform-as-a-service product that is marketed as a way to get your applications into the cloud without necessarily knowing all of the infrastructure bits and pieces to do so. Google App […]
Building Better Data Visualization Experiences: Part 2 of 2
If you don't have a Ph.D. in data science, the raw data might be difficult to comprehend. This is where data visualization comes in.
Unleashing Feature Flags onto Kafka Consumers
Feature flags are a tool to strategically enable or disable functionality at runtime. They are often used to drive different user experiences but can also be useful in real-time data systems. In this post, we’ll […]
A security model for developers
Software security is more important than ever, but developing secure applications is more confusing than ever. TLS, mTLS, RBAC, SAML, OAUTH, OWASP, GDPR, SASL, RSA, JWT, cookie, attack vector, DDoS, firewall, VPN, security groups, exploit, […]