From db2127a7271faa902ce7f9a988ee95d515ab9062 Mon Sep 17 00:00:00 2001
From: Lei Xia <lxia@linkedin.com>
Date: Tue, 21 Oct 2014 15:53:57 -0700
Subject: [PATCH] Add integration tests to verify that dropping a partition
 from resource ideal-state should bring partition to initial state and then
 DROPPED state (for AUTO, SEMI_AUTO, and CUSTOM modes)

---
 .../apache/helix/integration/TestDrop.java    | 33 ++++++++++++++-----
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/helix-core/src/test/java/org/apache/helix/integration/TestDrop.java b/helix-core/src/test/java/org/apache/helix/integration/TestDrop.java
index 63b566610a..7cfd22702a 100644
--- a/helix-core/src/test/java/org/apache/helix/integration/TestDrop.java
+++ b/helix-core/src/test/java/org/apache/helix/integration/TestDrop.java
@@ -19,6 +19,7 @@
  * under the License.
  */
 
+import java.util.ArrayList;
 import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
@@ -45,6 +46,7 @@
 import org.apache.helix.tools.ClusterStateVerifier;
 import org.apache.helix.tools.ClusterStateVerifier.BestPossAndExtViewZkVerifier;
 import org.testng.Assert;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 public class TestDrop extends ZkTestBase {
@@ -462,14 +464,27 @@ public void testDropSchemataResource() throws Exception {
     System.out.println("END " + clusterName + " at " + new Date(System.currentTimeMillis()));
   }
 
-  /**
-   * Drop a single partition in a resource of semi-auto mode
-   */
-  @Test
-  public void testDropSinglePartitionSemiAuto() throws Exception {
-    String className = TestHelper.getTestClassName();
-    String methodName = TestHelper.getTestMethodName();
-    String clusterName = className + "_" + methodName;
+  @DataProvider(name = "RebalanceModeProvider")
+  public static Object [][] rebalanceModes() {
+    IdealState.RebalanceMode modes [] = IdealState.RebalanceMode.values();
+    Object [][] args = new Object[modes.length][];
+    for (int i = 0; i < modes.length; i++) {
+      args[i] = new Object [] { modes[i] };
+    }
+    return args;
+  }
+
+
+  @Test(dataProvider = "RebalanceModeProvider")
+  public void testDropSinglePartition(IdealState.RebalanceMode mode)
+      throws Exception {
+    if (!mode.equals(IdealState.RebalanceMode.FULL_AUTO) &&
+        !mode.equals(IdealState.RebalanceMode.SEMI_AUTO) &&
+        !mode.equals(IdealState.RebalanceMode.CUSTOMIZED)) {
+      return;
+    }
+
+    String clusterName = TestHelper.getTestClassName() + "_" + TestHelper.getTestMethodName() + "_" + mode.name();
     int n = 2;
 
     System.out.println("START " + clusterName + " at " + new Date(System.currentTimeMillis()));
@@ -481,7 +496,7 @@ public void testDropSinglePartitionSemiAuto() throws Exception {
         4, // partitions per resource
         n, // number of nodes
         2, // replicas
-        "MasterSlave", true); // do rebalance
+        "MasterSlave", mode, true); // do rebalance
 
     ZKHelixDataAccessor accessor = new ZKHelixDataAccessor(clusterName, _baseAccessor);
     PropertyKey.Builder keyBuilder = accessor.keyBuilder();