How to mock DriverManager.getConnection(…)?

问题:

I have a class, which connects to an H2 database and runs several SQL statements.

public class H2Persistence implements IPersistence {

    private Connection conn;

    @Override
    public void open() {
        try
        {
            Class.forName("org.h2.Driver");
            conn = DriverManager.getConnection(CONN_TYPE_USER_HOME);

            final Statement stmt = conn.createStatement();

            stmt.executeUpdate("CREATE TABLE PERSON(" +
                    "ID BIGINT,"+
                    "AGEGROUP VARCHAR(255),"+
                    "MONTHLY_INCOME_LEVEL VARCHAR(255)," +
                    "GENDER VARCHAR(1),"+
                    "HOUSEHOLD_ID BIGINT)");

        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }
...
}

I want to write a unit test, which verifies, that in the open method a certain SQL statement (DROP TABLE IF EXISTS PERSON) is executed.

In order to do this, I wrote following test:

import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.powermock.api.mockito.PowerMockito.mockStatic;
import static org.powermock.api.mockito.PowerMockito.when;


@RunWith(PowerMockRunner.class)
@PrepareForTest(DriverManager.class)
public class H2PersistenceTest {
    @Test
    public void testDropPersonIsCalled() throws SQLException {
        final Statement statement = mock(Statement.class);

        final Connection connection = mock(Connection.class);

        when(connection.createStatement()).thenReturn(statement);

        mockStatic(DriverManager.class);

        when(DriverManager.getConnection(H2Persistence.CONN_TYPE_USER_HOME)).thenReturn
                (connection);


        final H2Persistence objectUnderTest = new H2Persistence();

        objectUnderTest.open();
        verify(statement.executeUpdate("DROP TABLE IF EXISTS PERSON"));
    }
}

But it doesn't work - instead of the mock connection, DriverManager returns real connection.

How can I fix it (make DriverManager return connection mock in the test) ?

Here's the pom.xml of my project, maybe something is wrong there.

<?xml version="1.0" encoding="UTF-8"?>

<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>ru.mycompany</groupId>
    <artifactId>myproduct</artifactId>
    <version>1.0-SNAPSHOT</version>
    <properties>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <powermock.version>1.5.1</powermock.version>
        <maven.compiler.source>1.6</maven.compiler.source>
        <maven.compiler.target>1.6</maven.compiler.target>
    </properties>

    <dependencies>
        <dependency>
            <groupId>junit</groupId>
            <artifactId>junit</artifactId>
            <version>4.10</version>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>org.easytesting</groupId>
            <artifactId>fest-util</artifactId>
            <version>1.2.3</version>
        </dependency>
        <dependency>
            <groupId>org.easytesting</groupId>
            <artifactId>fest-assert-core</artifactId>
            <version>2.0M8</version>
        </dependency>
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>15.0</version>
        </dependency>
        <dependency>
            <groupId>org.mockito</groupId>
            <artifactId>mockito-all</artifactId>
            <version>1.9.5</version>
        </dependency>
        <dependency>
            <groupId>com.h2database</groupId>
            <artifactId>h2</artifactId>
            <version>1.3.173</version>
        </dependency>
        <dependency>
            <groupId>org.powermock</groupId>
            <artifactId>powermock-module-junit4</artifactId>
            <version>${powermock.version}</version>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>org.powermock</groupId>
            <artifactId>powermock-api-mockito</artifactId>
            <version>${powermock.version}</version>
            <scope>test</scope>
        </dependency>
    </dependencies>

</project>

回答1:

This one works (pay attention to the imports):

import static org.easymock.EasyMock.expect;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.powermock.api.easymock.PowerMock.mockStatic;
import static org.powermock.api.easymock.PowerMock.replay;


@RunWith(PowerMockRunner.class)
@PrepareForTest({DriverManager.class, H2Persistence.class})
public class H2PersistenceTest {
    @Test
    public void testDropPersonIsCalled() throws SQLException {
        final Statement statement = mock(Statement.class);

        final Connection connection = mock(Connection.class);

        when(connection.createStatement()).thenReturn(statement);

        mockStatic(DriverManager.class);

        expect(DriverManager.getConnection(H2Persistence.CONN_TYPE_USER_HOME))
                .andReturn(connection);
        expect(DriverManager.getConnection(null))
                .andReturn(null);

        replay(DriverManager.class);
        final H2Persistence objectUnderTest = new H2Persistence();

        objectUnderTest.open();

        verify(statement).executeUpdate("DROP TABLE IF EXISTS PERSON");
        verify(statement).executeUpdate(H2Persistence.CREATE_TABLE_PERSON);
    }
}

回答2:

The usual way to do this would be to factor out the connection creation into another class, and inject an instance of that into the class in question. You can then mock that new class.

In your case, something like this:

public class H2Persistence implements IPersistence {
    private final ConnectionFactory connectionFactory;
    private Connection conn;

    public H2Persistence(ConnectionFactory connectionFactory) {
        this.connectionFactory = connectionFactory;
    }

    @Override
    public void open() {
        try {
            conn = connectionFactory.createConnection(CONN_TYPE_USER_HOME);
            // etc
        }
        catch (ClassNotFoundException e) {
            e.printStackTrace();
        }
        catch (SQLException e) {
            e.printStackTrace();
        }
    }
}

public class ConnectionFactory {

    Connection createConnection(String connType) throws SQLException, ClassNotFoundException {
        Class.forName("org.h2.Driver");
        return DriverManager.getConnection(connType);
    }

}

In this particular case, even better would probably be to use the standard JDBC interface DataSource instead of your own connection factory class:

public class H2Persistence implements IPersistence {
    private final DataSource dataSource;
    private Connection conn;

    public H2Persistence(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    @Override
    public void open() {
        try {
            conn = dataSource.getConnection();
            // etc
        }
        catch (SQLException e) {
            e.printStackTrace();
        }
    }
}