From fa2028074af6252c4af1d5fb862f6c0c592564d8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B8=A9=E5=B0=91?= <szujobs@hotmail.com>
Date: Tue, 20 Mar 2012 02:29:32 +0800
Subject: [PATCH] enhance mock driver

---
 .../com/alibaba/druid/mock/MockDriver.java    | 32 +++++++++++++++++--
 .../alibaba/druid/mock/MockDriverMBean.java   | 28 ++++++++++++++++
 .../alibaba/druid/pool/DruidDataSource.java   |  7 +++-
 3 files changed, 64 insertions(+), 3 deletions(-)
 create mode 100644 src/main/java/com/alibaba/druid/mock/MockDriverMBean.java

diff --git a/src/main/java/com/alibaba/druid/mock/MockDriver.java b/src/main/java/com/alibaba/druid/mock/MockDriver.java
index c45b31bac0..c75cb72bd3 100644
--- a/src/main/java/com/alibaba/druid/mock/MockDriver.java
+++ b/src/main/java/com/alibaba/druid/mock/MockDriver.java
@@ -15,6 +15,7 @@
  */
 package com.alibaba.druid.mock;
 
+import java.lang.management.ManagementFactory;
 import java.sql.Blob;
 import java.sql.Clob;
 import java.sql.Connection;
@@ -32,12 +33,15 @@ import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.logging.Logger;
 
+import javax.management.MBeanServer;
+import javax.management.ObjectName;
+
 import com.alibaba.druid.logging.Log;
 import com.alibaba.druid.logging.LogFactory;
 import com.alibaba.druid.mock.handler.MockExecuteHandler;
 import com.alibaba.druid.mock.handler.MySqlMockExecuteHandlerImpl;
 
-public class MockDriver implements Driver {
+public class MockDriver implements Driver, MockDriverMBean {
 
     private final static Log               LOG                   = LogFactory.getLog(MockDriver.class);
 
@@ -61,6 +65,8 @@ public class MockDriver implements Driver {
 
     private boolean                        logExecuteQueryEnable = true;
 
+    private final static String            MBEAN_NAME            = "com.alibaba.druid:type=MockDriver";
+
     static {
         registerDriver(instance);
     }
@@ -85,6 +91,17 @@ public class MockDriver implements Driver {
         return connectionIdSeed.incrementAndGet();
     }
 
+    public void closeAllConnections() throws SQLException {
+        for (int i = 0, size = this.connections.size(); i < size; ++i) {
+            Connection conn = this.connections.get(size - i - 1);
+            conn.close();
+        }
+    }
+    
+    public int getConnectionsSize() {
+        return this.connections.size();
+    }
+    
     public List<MockConnection> getConnections() {
         return connections;
     }
@@ -111,6 +128,17 @@ public class MockDriver implements Driver {
         try {
             DriverManager.registerDriver(driver);
 
+            try {
+                MBeanServer mbeanServer = ManagementFactory.getPlatformMBeanServer();
+
+                ObjectName objectName = new ObjectName(MBEAN_NAME);
+                if (!mbeanServer.isRegistered(objectName)) {
+                    mbeanServer.registerMBean(instance, objectName);
+                }
+            } catch (Exception ex) {
+                LOG.error("register druid-driver mbean error", ex);
+            }
+
             return true;
         } catch (Exception e) {
             LOG.error("registerDriver error", e);
@@ -126,7 +154,7 @@ public class MockDriver implements Driver {
     public void setExecuteHandler(MockExecuteHandler executeHandler) {
         this.executeHandler = executeHandler;
     }
-
+    
     @Override
     public Connection connect(String url, Properties info) throws SQLException {
         if (!acceptsURL(url)) {
diff --git a/src/main/java/com/alibaba/druid/mock/MockDriverMBean.java b/src/main/java/com/alibaba/druid/mock/MockDriverMBean.java
new file mode 100644
index 0000000000..6a96c9fc7d
--- /dev/null
+++ b/src/main/java/com/alibaba/druid/mock/MockDriverMBean.java
@@ -0,0 +1,28 @@
+package com.alibaba.druid.mock;
+
+import java.sql.SQLException;
+
+
+public interface MockDriverMBean {
+    long getConnectionCloseCount();
+    
+    int getMajorVersion();
+    
+    int getMinorVersion();
+    
+    boolean jdbcCompliant();
+    
+    boolean acceptsURL(String url) throws SQLException;
+    
+    boolean isLogExecuteQueryEnable();
+    
+    void setLogExecuteQueryEnable(boolean logExecuteQueryEnable);
+    
+    long getIdleTimeCount();
+    
+    void setIdleTimeCount(long idleTimeCount);
+    
+    void closeAllConnections() throws SQLException;
+    
+    int getConnectionsSize();
+}
diff --git a/src/main/java/com/alibaba/druid/pool/DruidDataSource.java b/src/main/java/com/alibaba/druid/pool/DruidDataSource.java
index dfc4f48a41..ad0d72bb9b 100644
--- a/src/main/java/com/alibaba/druid/pool/DruidDataSource.java
+++ b/src/main/java/com/alibaba/druid/pool/DruidDataSource.java
@@ -41,6 +41,7 @@ import com.alibaba.druid.VERSION;
 import com.alibaba.druid.filter.Filter;
 import com.alibaba.druid.logging.Log;
 import com.alibaba.druid.logging.LogFactory;
+import com.alibaba.druid.mock.MockDriver;
 import com.alibaba.druid.pool.DruidPooledPreparedStatement.PreparedStatementKey;
 import com.alibaba.druid.pool.vendor.InformixExceptionSorter;
 import com.alibaba.druid.pool.vendor.MSSQLValidConnectionChecker;
@@ -276,7 +277,11 @@ public class DruidDataSource extends DruidAbstractDataSource implements DruidDat
                     this.driverClass = JdbcUtils.getDriverClassName(this.jdbcUrl);
                 }
 
-                driver = JdbcUtils.createDriver(driverClass);
+                if (MockDriver.class.getName().equals(driverClass)) {
+                    driver = MockDriver.instance;
+                } else {
+                    driver = JdbcUtils.createDriver(driverClass);
+                }
             } else {
                 if (this.driverClass == null) {
                     this.driverClass = driver.getClass().getName();