SqlStatementCountUtils.java

import ch.qos.logback.classic.Level;
import ch.qos.logback.classic.Logger;
import ch.qos.logback.classic.spi.ILoggingEvent;
import ch.qos.logback.core.read.ListAppender;
import jakarta.annotation.Nonnull;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.function.Supplier;

import static org.assertj.core.api.Assertions.assertThat;

public class SqlStatementCountUtils {

    @Nonnull
    private final List<ILoggingEvent> sqlStatementsList;

    public SqlStatementCountUtils() {
        final ListAppender<ILoggingEvent> listAppender = new ListAppender<>();
        listAppender.start();
        final Logger logger = (Logger) LoggerFactory.getLogger("org.hibernate.SQL");

        logger.setLevel(Level.DEBUG);
        logger.addAppender(listAppender);

        sqlStatementsList = listAppender.list;
    }

    public <T> T assertSqlStatementsCount(int expectedCount, @Nonnull final Supplier<T> supplier) {
        final var returnValue = supplier.get();
        assertSizeAndClear(expectedCount);
        return returnValue;
    }

    public void assertSqlStatementsCount(int expectedCount, @Nonnull final Runnable runnable) {
        runnable.run();
        assertSizeAndClear(expectedCount);
    }

    private void assertSizeAndClear(int expectedCount) {
        assertThat(sqlStatementsList).hasSize(expectedCount);
        sqlStatementsList.clear();
    }

}