From 9a02c925b424cead5d38d805a14cf4372597de05 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B8=A9=E5=B0=91?= <szujobs@hotmail.com>
Date: Wed, 6 Jun 2012 08:59:34 +0800
Subject: [PATCH] mock refactor

---
 .../com/alibaba/druid/mock/MockConnection.java   | 16 +++++++++++++---
 .../java/com/alibaba/druid/mock/MockDriver.java  |  7 ++++++-
 .../druid/mock/MockPreparedStatement.java        | 10 ++++++----
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/alibaba/druid/mock/MockConnection.java b/src/main/java/com/alibaba/druid/mock/MockConnection.java
index 8ded645443..51d4d8e82a 100644
--- a/src/main/java/com/alibaba/druid/mock/MockConnection.java
+++ b/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) {
diff --git a/src/main/java/com/alibaba/druid/mock/MockDriver.java b/src/main/java/com/alibaba/druid/mock/MockDriver.java
index 5cdbff2de5..d7f6f7cc6e 100644
--- a/src/main/java/com/alibaba/druid/mock/MockDriver.java
+++ b/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);
+    }
 }
diff --git a/src/main/java/com/alibaba/druid/mock/MockPreparedStatement.java b/src/main/java/com/alibaba/druid/mock/MockPreparedStatement.java
index bfa0df5140..6159a3693c 100644
--- a/src/main/java/com/alibaba/druid/mock/MockPreparedStatement.java
+++ b/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);
     }