package io.dropwizard.db;

import com.codahale.metrics.MetricRegistry;
import io.dropwizard.util.Duration;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;

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

public class DataSourceFactoryTest {
    private final MetricRegistry metricRegistry = new MetricRegistry();
    private final DataSourceFactory factory = new DataSourceFactory();

    private ManagedDataSource dataSource;

    @Before
    public void setUp() throws Exception {
        factory.setUrl("jdbc:h2:mem:DbTest-" + System.currentTimeMillis());
        factory.setUser("sa");
        factory.setDriverClass("org.h2.Driver");
        factory.setValidationQuery("SELECT 1");
    }

    @After
    public void tearDown() throws Exception {
        dataSource.stop();
    }

    private ManagedDataSource dataSource() throws Exception {
        dataSource = factory.build(metricRegistry, "test");
        dataSource.start();
        return dataSource;
    }

    @Test
    public void buildsAConnectionPoolToTheDatabase() throws Exception {
        try (Connection connection = dataSource().getConnection()) {
            try (PreparedStatement statement = connection.prepareStatement(
                    "select 1")) {
                try (ResultSet set = statement.executeQuery()) {
                    while (set.next()) {
                        assertThat(set.getInt(1)).isEqualTo(1);
                    }
                }
            }
        }
    }

    @Test
    public void testNoValidationQueryTimeout() throws Exception {
        try (Connection connection = dataSource().getConnection()) {
            try (PreparedStatement statement = connection.prepareStatement(
                    "select 1")) {
                assertThat(statement.getQueryTimeout()).isEqualTo(0);
            }
        }
    }

    @Test
    public void testValidationQueryTimeoutIsSet() throws Exception {
        factory.setValidationQueryTimeout(Duration.seconds(3));

        try (Connection connection = dataSource().getConnection()) {
            try (PreparedStatement statement = connection.prepareStatement(
                    "select 1")) {
                assertThat(statement.getQueryTimeout()).isEqualTo(3);
            }
        }
    }

}

@@ -262,16 +262,46 @@
 sultSet;
+%0Aimport java.sql.SQLException;
 %0A%0Aimport
@@ -458,24 +458,25 @@
 Registry();%0A
+%0A
     private 
@@ -475,22 +475,16 @@
 private 
-final 
 DataSour
@@ -504,36 +504,9 @@
 tory
- = new DataSourceFactory();%0A
+;
 %0A   
@@ -581,34 +581,60 @@
 setUp() 
-throws Exception %7B
+%7B%0A        factory = new DataSourceFactory();
 %0A       
@@ -886,32 +886,70 @@
 ows Exception %7B%0A
+        if (null != dataSource) %7B%0A    
         dataSour
@@ -955,24 +955,34 @@
 rce.stop();%0A
+        %7D%0A
     %7D%0A%0A    p
@@ -1386,37 +1386,16 @@
 atement(
-%0A                    
 %22select 
@@ -1862,37 +1862,16 @@
 atement(
-%0A                    
 %22select 
@@ -2266,37 +2266,16 @@
 atement(
-%0A                    
 %22select 
@@ -2382,10 +2382,331 @@
 %0A    %7D%0A%0A
+    @Test(expected = SQLException.class)%0A    public void invalidJDBCDriverClassThrowsSQLException() throws SQLException %7B%0A        final DataSourceFactory factory = new DataSourceFactory();%0A        factory.setDriverClass(%22org.example.no.driver.here%22);%0A%0A        factory.build(metricRegistry, %22test%22).getConnection();%0A    %7D%0A
 %7D%0A

Add test for handling invalid JDBC driver class in DataSourceFactory