ha

src/test/java/com/alibaba/druid/bvt/pool/ha/CobarDataSourceTest.java

@@ -4,14 +4,17 @@ import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
-
-import org.junit.Assert;
+import java.util.concurrent.atomic.AtomicBoolean;
 
 import junit.framework.TestCase;
 
+import org.junit.Assert;
+
 import com.alibaba.druid.mock.MockConnection;
+import com.alibaba.druid.pool.DruidDataSource;
 import com.alibaba.druid.pool.ha.cobar.CobarConfigLoader;
 import com.alibaba.druid.pool.ha.cobar.CobarDataSource;
+import com.alibaba.druid.pool.ha.cobar.CobarFailureDetecter;
 import com.alibaba.fastjson.JSONArray;
 import com.alibaba.fastjson.JSONObject;
 
@@ -40,8 +43,11 @@ public class CobarDataSourceTest extends TestCase {
 
         cobarList.add(cobarB);
 
-        String mockUrlA = "jdbc:mock://mock:80/cobarA";
-        String mockUrlB = "jdbc:mock://mock:81/cobarB";
+        final String mockUrlA = "jdbc:mock://mock:80/cobarA";
+        final String mockUrlB = "jdbc:mock://mock:81/cobarB";
+        
+        final AtomicBoolean statusA = new AtomicBoolean(true);
+        final AtomicBoolean statusB = new AtomicBoolean(true);
 
         config.put("cobarList", cobarList);
 
@@ -50,7 +56,8 @@ public class CobarDataSourceTest extends TestCase {
         dataSource.setUsername("test");
         dataSource.setPassword("");
         dataSource.setMaxWait(1);
-        dataSource.setConfigLoadPeriodMillis(3);
+        dataSource.setConfigLoadPeriodMillis(10);
+        dataSource.setFailureDetectPeriodMillis(1);
 
         CobarConfigLoader configLoader = new CobarConfigLoader(dataSource) {
 
@@ -65,6 +72,23 @@ public class CobarDataSourceTest extends TestCase {
             }
         };
         dataSource.setConfigLoader(configLoader);
+        
+        CobarFailureDetecter failureDetector = new CobarFailureDetecter() {
+            public boolean isValidConnection(DruidDataSource dataSource, Connection conn) {
+                MockConnection mockConn;
+                try {
+                    mockConn = conn.unwrap(MockConnection.class);
+                } catch (SQLException e) {
+                    e.printStackTrace();
+                    return false;
+                }
+                if (mockConn.getUrl().equals(mockUrlA)) {
+                    return statusA.get();
+                }
+                return statusB.get();
+            }
+        };
+        dataSource.setFailureDetector(failureDetector);
 
         {
             Connection conn = dataSource.getConnection();
@@ -113,6 +137,55 @@ public class CobarDataSourceTest extends TestCase {
             stmt.close();
             conn.close();
         }
+        
+        cobarList.add(cobarA);
+        Assert.assertEquals(2, cobarList.size());
+        Thread.sleep(dataSource.getConfigLoadPeriodMillis() * 2);
+        for (int i = 0; i < 100; ++i) {
+            Connection conn = dataSource.getConnection();
+            Statement stmt = conn.createStatement();
+            ResultSet rs = stmt.executeQuery("select 1");
+
+            rs.next();
+            rs.close();
+            stmt.close();
+            conn.close();
+        }
+        
+        statusA.set(false);
+        Thread.sleep(dataSource.getFailureDetectPeriodMillis() * 2);
+        for (int i = 0; i < 100; ++i) {
+            Connection conn = dataSource.getConnection();
+            Statement stmt = conn.createStatement();
+            ResultSet rs = stmt.executeQuery("select 1");
+
+            MockConnection mockConn = conn.unwrap(MockConnection.class);
+            Assert.assertNotNull(mockConn);
+            Assert.assertEquals(mockUrlB, mockConn.getUrl());
+
+            rs.next();
+            rs.close();
+            stmt.close();
+            conn.close();
+        }
+        
+        statusA.set(true);
+        statusB.set(false);
+        Thread.sleep(dataSource.getFailureDetectPeriodMillis() * 2);
+        for (int i = 0; i < 100; ++i) {
+            Connection conn = dataSource.getConnection();
+            Statement stmt = conn.createStatement();
+            ResultSet rs = stmt.executeQuery("select 1");
+
+            MockConnection mockConn = conn.unwrap(MockConnection.class);
+            Assert.assertNotNull(mockConn);
+            Assert.assertEquals(mockUrlA, mockConn.getUrl());
+
+            rs.next();
+            rs.close();
+            stmt.close();
+            conn.close();
+        }
     }
 
 }