add more unit tests

ambry-clustermap/src/main/java/com.github.ambry.clustermap/AmbryStateModelDefinition.java

@@ -31,12 +31,12 @@
  *                  ---> BOOTSTRAP --->
  *
  */
-public class AmbryStateModelDefinition {
-  public static final String AMBRY_LEADER_STANDBY_MODEL = "AmbryLeaderStandby";
+class AmbryStateModelDefinition {
+  static final String AMBRY_LEADER_STANDBY_MODEL = "AmbryLeaderStandby";
 
   private static final String UPPER_BOUND_REPLICATION_FACTOR = "R";
 
-  public static StateModelDefinition getDefinition() {
+  static StateModelDefinition getDefinition() {
 
     StateModelDefinition.Builder builder = new StateModelDefinition.Builder(AMBRY_LEADER_STANDBY_MODEL);
     // Init state

ambry-clustermap/src/test/java/com.github.ambry.clustermap/AmbryStateModelFactoryTest.java

@@ -0,0 +1,62 @@
+/**
+ * Copyright 2019 LinkedIn Corp. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ */
+package com.github.ambry.clustermap;
+
+import com.github.ambry.config.ClusterMapConfig;
+import com.github.ambry.config.VerifiableProperties;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Properties;
+import org.apache.helix.participant.statemachine.StateModel;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+import static org.junit.Assert.*;
+
+
+/**
+ * Test for {@link AmbryStateModelFactory}
+ */
+@RunWith(Parameterized.class)
+public class AmbryStateModelFactoryTest {
+  private final ClusterMapConfig clustermapConfig;
+
+  @Parameterized.Parameters
+  public static List<Object[]> data() {
+    return Arrays.asList(
+        new Object[][]{{ClusterMapConfig.DEFAULT_STATE_MODEL_DEF}, {ClusterMapConfig.AMBRY_STATE_MODEL_DEF}});
+  }
+
+  public AmbryStateModelFactoryTest(String stateModelDef) {
+    Properties props = new Properties();
+    props.setProperty("clustermap.host.name", "localhost");
+    props.setProperty("clustermap.port", "2200");
+    props.setProperty("clustermap.cluster.name", "AmbryStateModelFactoryTest");
+    props.setProperty("clustermap.state.model.definition", stateModelDef);
+    props.setProperty("clustermap.datacenter.name", "DC0");
+    clustermapConfig = new ClusterMapConfig(new VerifiableProperties(props));
+  }
+
+  @Test
+  public void testDifferentStateModelDefs() {
+    AmbryStateModelFactory factory = new AmbryStateModelFactory(clustermapConfig);
+    StateModel stateModel = factory.createNewStateModel("0", "1");
+    if ((clustermapConfig.clustermapStateModelDefinition).equals(ClusterMapConfig.DEFAULT_STATE_MODEL_DEF)) {
+      assertTrue("Unexpected state model def", stateModel instanceof AmbryDefaultStateModel);
+    } else {
+      assertTrue("Unexpected state model def", stateModel instanceof AmbryPartitionStateModel);
+    }
+  }
+}

ambry-tools/src/integration-test/java/com.github.ambry/clustermap/HelixBootstrapUpgradeToolTest.java

@@ -163,6 +163,52 @@ public void testParseDcSet() throws Exception {
     Assert.assertEquals(expected, HelixBootstrapUpgradeUtil.parseAndUpdateDcInfoFromArg("all", zkLayoutPath).keySet());
   }
 
+  /**
+   * Test {@link HelixBootstrapUpgradeUtil} addStateModelDef() method.
+   */
+  @Test
+  public void testAddStateModelDef() throws Exception {
+    Utils.writeJsonObjectToFile(zkJson, zkLayoutPath);
+    Utils.writeJsonObjectToFile(testHardwareLayout.getHardwareLayout().toJSONObject(), hardwareLayoutPath);
+    Utils.writeJsonObjectToFile(testPartitionLayout.getPartitionLayout().toJSONObject(), partitionLayoutPath);
+    // test add state model to non-exist cluster, which should fail
+    try {
+      HelixBootstrapUpgradeUtil.addStateModelDef(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
+          CLUSTER_NAME_PREFIX, dcStr, DEFAULT_MAX_PARTITIONS_PER_RESOURCE, new HelixAdminFactory(),
+          ClusterMapConfig.AMBRY_STATE_MODEL_DEF);
+      fail("should fail due to non-exist cluster");
+    } catch (IllegalStateException e) {
+      // expected
+    }
+    // bootstrap a cluster
+    HelixBootstrapUpgradeUtil.bootstrapOrUpgrade(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
+        CLUSTER_NAME_PREFIX, dcStr, DEFAULT_MAX_PARTITIONS_PER_RESOURCE, false, false, new HelixAdminFactory(), true,
+        ClusterMapConfig.DEFAULT_STATE_MODEL_DEF);
+    // add new state model def
+    HelixBootstrapUpgradeUtil.addStateModelDef(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
+        CLUSTER_NAME_PREFIX, dcStr, DEFAULT_MAX_PARTITIONS_PER_RESOURCE, new HelixAdminFactory(),
+        ClusterMapConfig.AMBRY_STATE_MODEL_DEF);
+    // add existing state model def should be no-op
+    HelixBootstrapUpgradeUtil.addStateModelDef(hardwareLayoutPath, partitionLayoutPath, zkLayoutPath,
+        CLUSTER_NAME_PREFIX, dcStr, DEFAULT_MAX_PARTITIONS_PER_RESOURCE, new HelixAdminFactory(),
+        ClusterMapConfig.DEFAULT_STATE_MODEL_DEF);
+    // ensure that active dcs have new state model def
+    String clusterName = CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP;
+    for (Datacenter dc : testHardwareLayout.getHardwareLayout().getDatacenters()) {
+      ZkInfo zkInfo = dcsToZkInfo.get(dc.getName());
+      ZKHelixAdmin admin = new ZKHelixAdmin("localhost:" + zkInfo.getPort());
+      if (!activeDcSet.contains(dc.getName())) {
+        Assert.assertFalse("Cluster should not be present, as dc " + dc.getName() + " is not enabled",
+            admin.getClusters().contains(CLUSTER_NAME_PREFIX + CLUSTER_NAME_IN_STATIC_CLUSTER_MAP));
+      } else {
+        assertEquals("Mismatch in number of state model defs in cluster", 2,
+            admin.getStateModelDefs(clusterName).size());
+        assertTrue("Missing ambry state model in cluster",
+            admin.getStateModelDefs(clusterName).contains(ClusterMapConfig.AMBRY_STATE_MODEL_DEF));
+      }
+    }
+  }
+
   /**
    * Test the case where the zkHosts JSON does not have an entry for every Datacenter in the static clustermap.
    */