Added more zone expansion test for RebalanceBatchPlan

kthawan · Jun 20, 2013 · d344725 · d344725
1 parent d4c314f
commit d344725
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/test/common/voldemort/ClusterTestUtils.java b/test/common/voldemort/ClusterTestUtils.java
@@ -568,6 +568,26 @@ public static List<RebalancePartitionsInfo> getBatchPlan(Cluster currentCluster,
         return rebalancePlan.getBatchPlan();
     }
 
+    /**
+     * Given the current and target cluster metadata, along with your store
+     * definition, return the batch plan.
+     * 
+     * @param currentCluster Current cluster metadata
+     * @param targetCluster Target cluster metadata
+     * @param storeDef List of store definitions
+     * @return list of tasks for this batch plan
+     */
+    public static List<RebalancePartitionsInfo> getBatchPlan(Cluster currentCluster,
+                                                             List<StoreDefinition> currentStoreDefs,
+                                                             Cluster finalCluster,
+                                                             List<StoreDefinition> finalStoreDefs) {
+        RebalanceBatchPlan rebalancePlan = new RebalanceBatchPlan(currentCluster,
+                                                                  currentStoreDefs,
+                                                                  finalCluster,
+                                                                  finalStoreDefs);
+        return rebalancePlan.getBatchPlan();
+    }
+
     /**
      * Constructs a plan to rebalance from current state (cCluster/cStores) to
      * final state (fCluster/fStores). Uses default values for the planning.

diff --git a/test/unit/voldemort/client/rebalance/ZonedRebalanceBatchPlanTest.java b/test/unit/voldemort/client/rebalance/ZonedRebalanceBatchPlanTest.java
@@ -108,9 +108,20 @@ public void testClusterExpansion() {
     public void testZoneExpansion() {
         List<RebalancePartitionsInfo> batchPlan;
 
-        // Two-to-three zones
+        // Two-to-three zones (i)
         batchPlan = ClusterTestUtils.getBatchPlan(zzeZoneExpansion, zzzZoneExpansionXXP, zzzStores);
         assertTrue(batchPlan.size() > 0);
+
+        // Two-to-three zones (ii)
+        batchPlan = ClusterTestUtils.getBatchPlan(zzCurrent,
+                                                  zzStores,
+                                                  zzzZoneExpansionXXP,
+                                                  zzzStores);
+        assertTrue(batchPlan.size() > 0);
+
+        // Two-to-three zones (iii)
+        batchPlan = ClusterTestUtils.getBatchPlan(zzCurrent, zzStores, zzzCurrent, zzzStores);
+        assertTrue(batchPlan.size() > 0);
     }
 
 }