Spring `JpaRepository` to be used only in test scope

Question

I have a JpaRepository, that looks like

@Repository
public interface CustomRepository extends JpaRepository<EntityType, Integer> {
    // Methods
}

which could possibly have queries that would run for long, in which case I need to enforce a timeout. I have successfully added the timeout-related configuration (the connection pool being used is Druid, if that matters) and now I want to test it in a unit test. I am using the following method in my JpaRepository<T, ID> interface.

@Query(value = "SELECT benchmark(:count, MD5('3094803'))", nativeQuery = true)
void run(@Param("count") Long count);

This method runs successfully and demonstrates the expected behavior. However, given that this method will run longer as the value given to parameter count becomes larger, having this in the production code, just for the sake of testing the timeouts, bothers me to my core, as this might end up being a vulnerability that could be leveraged to launch a denial attack.

So to the question, is some way I can use this exact method in my test scope, without having that going in the production code?

Answer 1

Turns out, it's not that complicated. Given that this project is running on Spring, I can have an extension of the above repository in my test sources, as follows.

package com.project.package.repo;

import org.springframework.data.jpa.repository.Query;
import org.springframework.data.repository.query.Param;
import org.springframework.stereotype.Repository;

@Repository
public interface TimedCustomRepository extends CustomRepository {
    @Query(value = "SELECT benchmark(:count, MD5('3094803'))", nativeQuery = true)
    void run(@Param("count") Long count);
}

As I have JUnit4, I can have a test class which will run on Spring boot, like follows.

package om.project.package.repo;

import com.github.database.rider.spring.api.DBRider;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.orm.jpa.JpaSystemException;
import org.springframework.test.context.junit4.SpringRunner;

@SpringBootTest
@RunWith(SpringRunner.class)
@DBRider(dataSourceBeanName = "primaryDataSource") //Data source wiring; Not that important.
public class TransactionHistoryJpaRepoTest {

    @Autowired //---(1)
    private TimedCustomRepository timedCustomRepository;

    @Test(expected = JpaSystemException.class)
    public void whenQueryTimeoutOccurs() {
        timedCustomRepository.run(100000000L);
    }
}

The property at (1) will be wired using the repository bean which we have created above, and this test will execute as expected. Given the bean TimedCustomRepository extends CustomRepository, the data source configuration and everything will be the same. Most importantly, since this method with a long running query is now only on the test scope, it will not have any impact beyond the test scope.