JUnit @Rule for a SQL transaction

I needed to run some Java test code in a SQL transaction. It modifies the database, changes that should be 1) invisible to other tests, and 2) thrown away after tests finish.

I am using JUnit (of course) so I wrote a @Rule to provide a transaction to my test methods (note I'm using Lombok for convenience):

@RequiredArgsConstructor
public final class SQLTransactionRule
        extends ExternalResource {
    private final DataSource original;
    private final String user;
    private final String password;
    private final Consumer testDataSource;

    private Connection conn;

    public SQLTransactionRule(final DataSource original,
            final Consumer testDataSource) {
        this(original, null, null, testDataSource);
    }

    @Override
    protected void before()
            throws Throwable {
        super.before();
        beginTransaction();
        testDataSource.accept(testDataSource(conn));
    }

    @Override
    protected void after() {
        testDataSource.accept(null);
        rollbackTransaction();
        super.after();
    }

    private void beginTransaction() {
        try {
            final Connection conn = original.getConnection(user, password);
            conn.setAutoCommit(false);
            this.conn = conn;
        } catch (final SQLException e) {
            throw new RuntimeException(
                    format("Cannot create transaction from %s: %s", original,
                            getRootCause(e)), e);
        }
    }

    private void rollbackTransaction() {
        try {
            conn.rollback();
        } catch (final SQLException e) {
            throw new RuntimeException(
                    format("Cannot rollback transaction to %s: %s", original,
                            getRootCause(e)), e);
        }
    }

    private DataSource testDataSource(final Connection conn) {
        return (DataSource) newProxyInstance(getClass().getClassLoader(),
                new Class[]{DataSource.class}, (proxy, method, args) -> {
                    System.out.println("SQLTransactionRule.testDataSource");
                    if (Object.class.equals(method.getDeclaringClass()))
                        return method.invoke(proxy, args);
                    else if ("getConnection".equals(method.getName()))
                        return conn;
                    else
                        return method.invoke(original, args);
                });
    }
}

The motivation for the Consumer of "testDataSource" and for the proxied data source was avoidance of dependencies. Original code has this convenience method, but I decided that was a concern of the test code not the @Rule:

public JdbcTemplate newJdbcTemplate() {
    return new JdbcTemplate(testDataSource);
}

Of course there are tests for the @Rule as well! And this post on Mockito was handy.

@RunWith(MockitoJUnitRunner.class)
public final class SQLTransactionRuleTest {
    @Mock
    private DataSource original;
    @Mock
    private Connection conn;

    @Rule
    public final ExpectedException thrown = ExpectedException.none();

    private SQLTransactionRule rule;
    private DataSource testDataSource;

    @Before
    public void setUp()
            throws SQLException {
        when(original.getConnection(anyString(), anyString())).
                thenReturn(conn);
        rule = new SQLTransactionRule(original,
                testDataSource -> this.testDataSource = testDataSource);
    }

    @Test
    public void shouldExecuteStatement()
            throws Throwable {
        final boolean[] called = {false};
        executeStatement(() -> called[0] = true);
        assertThat(called[0], is(true));
    }

    @Test
    public void shouldPublishTestDataSource()
            throws Throwable {
        executeStatement(
                () -> assertThat(testDataSource, is(notNullValue())));
    }

    @Test
    public void shouldReuseOriginalConnection()
            throws Throwable {
        executeStatement(() -> assertThat(testDataSource.getConnection(),
                is(sameInstance(conn))));
    }

    @Test
    public void shouldTransact()
            throws Throwable {
        executeStatement(() -> verify(conn).setAutoCommit(eq(false)));
        verify(conn).rollback();
    }

    @Test
    public void shouldDescribe()
            throws Throwable {
        thrown.expect(AssertionError.class);
        thrown.expectMessage(containsString("alpha"));
        thrown.expectMessage(containsString("beta"));

        executeStatement(() -> assertThat("alpha", is(equalTo("beta"))));
    }

    @FunctionalInterface
    private interface RunMe {
        void execute()
                throws Throwable;
    }

    private void executeStatement(final RunMe statement)
            throws Throwable {
        rule.apply(new Statement() {
            @Override
            public void evaluate()
                    throws Throwable {
                statement.execute();
            }
        }, EMPTY).evaluate();
    }
}