mock refactor

src/main/java/com/alibaba/druid/mock/MockConnection.java

@@ -112,6 +112,9 @@ public class MockConnection extends ConnectionBase implements Connection {
     }
 
     private MockStatement createMockStatement() {
+        if (driver != null) {
+            return driver.createMockStatement(this);
+        }
         return new MockStatement(this);
     }
 
@@ -130,7 +133,7 @@ public class MockConnection extends ConnectionBase implements Connection {
             throw new MockConnectionClosedException();
         }
 
-        return new MockCallableStatement(this, sql);
+        return createMockCallableStatement(sql);
     }
 
     @Override
@@ -222,7 +225,7 @@ public class MockConnection extends ConnectionBase implements Connection {
             throw new MockConnectionClosedException();
         }
 
-        MockCallableStatement stmt = new MockCallableStatement(this, sql);
+        MockCallableStatement stmt = createMockCallableStatement(sql);
 
         stmt.setResultSetType(resultSetType);
         stmt.setResultSetConcurrency(resultSetConcurrency);
@@ -316,7 +319,7 @@ public class MockConnection extends ConnectionBase implements Connection {
             throw new MockConnectionClosedException();
         }
 
-        MockCallableStatement stmt = new MockCallableStatement(this, sql);
+        MockCallableStatement stmt = createMockCallableStatement(sql);
 
         stmt.setResultSetType(resultSetType);
         stmt.setResultSetConcurrency(resultSetConcurrency);
@@ -325,6 +328,13 @@ public class MockConnection extends ConnectionBase implements Connection {
         return stmt;
     }
 
+    private MockCallableStatement createMockCallableStatement(String sql) {
+        if (driver != null) {
+            return driver.createMockCallableStatement(this, sql);
+        }
+        return new MockCallableStatement(this, sql);
+    }
+
     @Override
     public PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException {
         if (closed) {

src/main/java/com/alibaba/druid/mock/MockDriver.java

@@ -27,6 +27,7 @@ import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLFeatureNotSupportedException;
 import java.sql.SQLXML;
+import java.sql.Statement;
 import java.util.List;
 import java.util.Properties;
 import java.util.concurrent.CopyOnWriteArrayList;
@@ -313,8 +314,12 @@ public class MockDriver implements Driver, MockDriverMBean {
     public MockPreparedStatement createMockPreparedStatement(MockConnection conn, String sql) {
         return new MockPreparedStatement(conn, sql);
     }
-    
+
     public MockStatement createMockStatement(MockConnection conn) {
         return new MockStatement(conn);
     }
+
+    public MockCallableStatement createMockCallableStatement(MockConnection conn, String sql) {
+        return new MockCallableStatement(conn, sql);
+    }
 }

src/main/java/com/alibaba/druid/mock/MockPreparedStatement.java

@@ -41,15 +41,17 @@ public class MockPreparedStatement extends PreparedStatementBase implements Prep
     @Override
     public ResultSet executeQuery() throws SQLException {
         checkOpen();
+        
+        MockConnection conn = getConnection();
 
-        if (getConnection() != null && getConnection().getDriver() != null) {
+        if (conn != null && conn.getDriver() != null) {
             return getConnection().getDriver().createResultSet(this);
         }
 
-        if (getConnection() != null) {
-            getConnection().handleSleep();
+        if (conn != null) {
+            conn.handleSleep();
         }
-
+        
         return new MockResultSet(this);
     }