helix-core/src/test/java/org/apache/helix/controller/rebalancer/waged/model/TestClusterModelProvider.java - helix - Git at Google

 package org.apache.helix.controller.rebalancer.waged.model;

 /*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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.  See the License for the
  * specific language governing permissions and limitations
  * under the License.
  */

 import java.io.IOException;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;

 import org.apache.helix.HelixConstants;
 import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
 import org.apache.helix.controller.rebalancer.waged.WagedRebalancer;
 import org.apache.helix.model.CurrentState;
 import org.apache.helix.model.IdealState;
 import org.apache.helix.model.InstanceConfig;
 import org.apache.helix.model.LiveInstance;
 import org.apache.helix.model.Partition;
 import org.apache.helix.model.Resource;
 import org.apache.helix.model.ResourceAssignment;
 import org.mockito.stubbing.Answer;
 import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;

 import static org.mockito.Matchers.anyString;
 import static org.mockito.Mockito.when;

 public class TestClusterModelProvider extends AbstractTestClusterModel {
   Set<String> _instances;

   @BeforeClass
   public void initialize() {
     super.initialize();
     _instances = new HashSet<>();
     _instances.add(_testInstanceId);
   }

   @Override
   protected ResourceControllerDataProvider setupClusterDataCache() throws IOException {
     ResourceControllerDataProvider testCache = super.setupClusterDataCache();

     // Set up mock idealstate
     Map<String, IdealState> isMap = new HashMap<>();
     for (String resource : _resourceNames) {
       IdealState is = new IdealState(resource);
       is.setNumPartitions(_partitionNames.size());
       is.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
       is.setStateModelDefRef("MasterSlave");
       is.setReplicas("3");
       is.setRebalancerClassName(WagedRebalancer.class.getName());
       _partitionNames.stream()
           .forEach(partition -> is.setPreferenceList(partition, Collections.emptyList()));
       isMap.put(resource, is);
     }
     when(testCache.getIdealState(anyString())).thenAnswer(
         (Answer<IdealState>) invocationOnMock -> isMap.get(invocationOnMock.getArguments()[0]));

     // Set up 2 more instances
     for (int i = 1; i < 3; i++) {
       String instanceName = _testInstanceId + i;
       _instances.add(instanceName);
       // 1. Set up the default instance information with capacity configuration.
       InstanceConfig testInstanceConfig = createMockInstanceConfig(instanceName);
       Map<String, InstanceConfig> instanceConfigMap = testCache.getInstanceConfigMap();
       instanceConfigMap.put(instanceName, testInstanceConfig);
       when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
       // 2. Mock the live instance node for the default instance.
       LiveInstance testLiveInstance = createMockLiveInstance(instanceName);
       Map<String, LiveInstance> liveInstanceMap = testCache.getLiveInstances();
       liveInstanceMap.put(instanceName, testLiveInstance);
       when(testCache.getLiveInstances()).thenReturn(liveInstanceMap);
     }

     return testCache;
   }

   @Test
   public void testGenerateClusterModel() throws IOException {
     ResourceControllerDataProvider testCache = setupClusterDataCache();
     // 1. test generating a cluster model with empty assignment
     ClusterModel clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances, Collections.emptyMap(), Collections.emptyMap());
     // There should be no existing assignment.
     Assert.assertFalse(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .anyMatch(resourceMap -> !resourceMap.isEmpty()));
     Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
         .anyMatch(node -> node.getAssignedReplicaCount() != 0));
     // Have all 3 instances
     Assert.assertEquals(
         clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
             .collect(Collectors.toSet()), _instances);
     // Shall have 2 resources and 4 replicas, since all nodes are in the same fault zone.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 4));

     // Adjust instance fault zone, so they have different fault zones.
     testCache.getInstanceConfigMap().values().stream()
         .forEach(config -> config.setZoneId(config.getInstanceName()));
     clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances, Collections.emptyMap(), Collections.emptyMap());
     // Shall have 2 resources and 12 replicas after fault zone adjusted.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 12));

     // 2. test with only one active node
     clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         Collections.singleton(_testInstanceId), Collections.emptyMap(), Collections.emptyMap());
     // Have only one instance
     Assert.assertEquals(
         clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
             .collect(Collectors.toSet()), Collections.singleton(_testInstanceId));
     // Shall have 4 assignable replicas because there is only one valid node.
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 4));

     // 3. test with no active instance
     clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         Collections.emptySet(), Collections.emptyMap(), Collections.emptyMap());
     // Have only one instance
     Assert.assertEquals(clusterModel.getAssignableNodes().size(), 0);
     // Shall have 0 assignable replicas because there are 0 valid nodes.
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.isEmpty()));

     // 4. test with baseline assignment
     // Mock a baseline assignment based on the current states.
     Map<String, ResourceAssignment> baselineAssignment = new HashMap<>();
     for (String resource : _resourceNames) {
       // <partition, <instance, state>>
       Map<String, Map<String, String>> assignmentMap = new HashMap<>();
       CurrentState cs = testCache.getCurrentState(_testInstanceId, _sessionId).get(resource);
       if (cs != null) {
         for (Map.Entry<String, String> stateEntry : cs.getPartitionStateMap().entrySet()) {
           assignmentMap.computeIfAbsent(stateEntry.getKey(), k -> new HashMap<>())
               .put(_testInstanceId, stateEntry.getValue());
         }
         ResourceAssignment assignment = new ResourceAssignment(resource);
         assignmentMap.keySet().stream().forEach(partition -> assignment
             .addReplicaMap(new Partition(partition), assignmentMap.get(partition)));
         baselineAssignment.put(resource, assignment);
       }
     }

     // Generate a cluster model based on the best possible assignment
     clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances, Collections.emptyMap(), baselineAssignment);
     // There should be 4 existing assignments in total (each resource has 2) in the specified instance
     Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .allMatch(resourceMap -> resourceMap.values().stream()
             .allMatch(partitionSet -> partitionSet.size() == 2)));
     Assert.assertEquals(
         clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 4);
     // Since each resource has 2 replicas assigned, the assignable replica count should be 10.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 10));

     // 5. test with best possible assignment but cluster topology is changed
     clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances,
         Collections.singletonMap(HelixConstants.ChangeType.CLUSTER_CONFIG, Collections.emptySet()),
         baselineAssignment);
     // There should be no existing assignment since the topology change invalidates all existing assignment
     Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .allMatch(resourceMap -> resourceMap.isEmpty()));
     Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
         .anyMatch(node -> node.getAssignedReplicaCount() != 0));
     // Shall have 2 resources and 12 replicas
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 12));

     // 6. test with best possible assignment and one resource config change
     // Generate a cluster model based on the same best possible assignment, but resource1 config is changed
     String changedResourceName = _resourceNames.get(0);
     clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances, Collections.singletonMap(HelixConstants.ChangeType.RESOURCE_CONFIG,
             Collections.singleton(changedResourceName)), baselineAssignment);
     // There should be no existing assignment for all the resource except for resource2
     Assert.assertEquals(clusterModel.getContext().getAssignmentForFaultZoneMap().size(), 1);
     Map<String, Set<String>> resourceAssignmentMap =
         clusterModel.getContext().getAssignmentForFaultZoneMap().get(_testInstanceId);
     // Should be only resource2 in the map
     Assert.assertEquals(resourceAssignmentMap.size(), 1);
     for (String resource : _resourceNames) {
       Assert
           .assertEquals(resourceAssignmentMap.getOrDefault(resource, Collections.emptySet()).size(),
               resource.equals(changedResourceName) ? 0 : 2);
     }
     // Only the first instance will have 2 assignment from resource2.
     for (String instance : _instances) {
       Assert.assertEquals(clusterModel.getAssignableNodes().get(instance).getAssignedReplicaCount(),
           instance.equals(_testInstanceId) ? 2 : 0);
     }
     // Shall have 2 resources and 12 replicas
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().keySet().size(), 2);
     for (String resource : _resourceNames) {
       Assert.assertEquals(clusterModel.getAssignableReplicaMap().get(resource).size(),
           resource.equals(changedResourceName) ? 12 : 10);
     }

     // 7. test with best possible assignment but the instance becomes inactive
     // Generate a cluster model based on the best possible assignment, but the assigned node is disabled
     Set<String> limitedActiveInstances = new HashSet<>(_instances);
     limitedActiveInstances.remove(_testInstanceId);
     clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         limitedActiveInstances, Collections.emptyMap(), baselineAssignment);
     // There should be no existing assignment.
     Assert.assertFalse(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .anyMatch(resourceMap -> !resourceMap.isEmpty()));
     Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
         .anyMatch(node -> node.getAssignedReplicaCount() != 0));
     // Have only 2 instances
     Assert.assertEquals(
         clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
             .collect(Collectors.toSet()), limitedActiveInstances);
     // Since only 2 instances are active, we shall have 8 assignable replicas in each resource.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 8));
   }

   @Test (dependsOnMethods = "testGenerateClusterModel")
   public void testGenerateClusterModelForPartialRebalance() throws IOException {
     ResourceControllerDataProvider testCache = setupClusterDataCache();
     // 1. test generating a cluster model with empty assignment
     ClusterModel clusterModel = ClusterModelProvider
         .generateClusterModelForPartialRebalance(testCache, _resourceNames.stream()
                 .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
             _instances, Collections.emptyMap(), Collections.emptyMap());
     // There should be no existing assignment.
     Assert.assertFalse(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .anyMatch(resourceMap -> !resourceMap.isEmpty()));
     Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
         .anyMatch(node -> node.getAssignedReplicaCount() != 0));
     // Have all 3 instances
     Assert.assertEquals(
         clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
             .collect(Collectors.toSet()), _instances);
     // Shall have 0 resources and 0 replicas since the baseline is empty. The partial rebalance
     // should not rebalance any replica.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 0);

     // Adjust instance fault zone, so they have different fault zones.
     testCache.getInstanceConfigMap().values().stream()
         .forEach(config -> config.setZoneId(config.getInstanceName()));

     // 2. test with a pair of identical best possible assignment and baseline assignment
     // Mock a best possible assignment based on the current states.
     Map<String, ResourceAssignment> bestPossibleAssignment = new HashMap<>();
     for (String resource : _resourceNames) {
       // <partition, <instance, state>>
       Map<String, Map<String, String>> assignmentMap = new HashMap<>();
       CurrentState cs = testCache.getCurrentState(_testInstanceId, _sessionId).get(resource);
       if (cs != null) {
         for (Map.Entry<String, String> stateEntry : cs.getPartitionStateMap().entrySet()) {
           assignmentMap.computeIfAbsent(stateEntry.getKey(), k -> new HashMap<>())
               .put(_testInstanceId, stateEntry.getValue());
         }
         ResourceAssignment assignment = new ResourceAssignment(resource);
         assignmentMap.keySet().stream().forEach(partition -> assignment
             .addReplicaMap(new Partition(partition), assignmentMap.get(partition)));
         bestPossibleAssignment.put(resource, assignment);
       }
     }
     Map<String, ResourceAssignment> baseline = new HashMap<>(bestPossibleAssignment);
     // Generate a cluster model for partial rebalance
     clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances, baseline, bestPossibleAssignment);
     // There should be 4 existing assignments in total (each resource has 2) in the specified instance
     Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .allMatch(resourceMap -> resourceMap.values().stream()
             .allMatch(partitionSet -> partitionSet.size() == 2)));
     Assert.assertEquals(
         clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 4);
     // Since the best possible matches the baseline, no replica needs to be reassigned.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 0);

     // 3. test with inactive instance in the baseline and the best possible assignment
     Set<String> partialInstanceList = new HashSet<>(_instances);
     partialInstanceList.remove(_testInstanceId);
     clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         partialInstanceList, baseline, bestPossibleAssignment);
     // Have the other 2 active instances
     Assert.assertEquals(clusterModel.getAssignableNodes().size(), 2);
     // All the replicas in the existing assignment should be rebalanced.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 2));
     // Shall have 0 assigned replicas
     Assert.assertTrue(clusterModel.getAssignableNodes().values().stream()
         .allMatch(assignableNode -> assignableNode.getAssignedReplicaCount() == 0));

     // 4. test with one resource that is only in the baseline
     String resourceInBaselineOnly = _resourceNames.get(0);
     Map<String, ResourceAssignment> partialBestPossibleAssignment =
         new HashMap<>(bestPossibleAssignment);
     partialBestPossibleAssignment.remove(resourceInBaselineOnly);
     // Generate a cluster mode with the adjusted best possible assignment
     clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances, baseline, partialBestPossibleAssignment);
     // There should be 2 existing assignments in total in the specified instance
     Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .allMatch(resourceMap -> resourceMap.values().stream()
             .allMatch(partitionSet -> partitionSet.size() == 2)));
     // Only the replicas of one resource require rebalance
     Assert.assertEquals(
         clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 2);
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 1);
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().containsKey(resourceInBaselineOnly));
     Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
         .allMatch(replicaSet -> replicaSet.size() == 2));

     // 5. test with one resource only in the best possible assignment
     String resourceInBestPossibleOnly = _resourceNames.get(1);
     Map<String, ResourceAssignment> partialBaseline = new HashMap<>(baseline);
     partialBaseline.remove(resourceInBestPossibleOnly);
     // Generate a cluster model with the adjusted baseline
     clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
         _resourceNames.stream()
             .collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
         _instances, partialBaseline, bestPossibleAssignment);
     // There should be 2 existing assignments in total and all of them require rebalance.
     Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
         .allMatch(resourceMap -> resourceMap.values().stream()
             .allMatch(partitionSet -> partitionSet.size() == 2)));
     Assert.assertEquals(
         clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 2);
     // No need to rebalance the replicas that are not in the baseline yet.
     Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 0);
   }
 }
	package org.apache.helix.controller.rebalancer.waged.model;

	/*
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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. See the License for the
	* specific language governing permissions and limitations
	* under the License.
	*/

	import java.io.IOException;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.Map;
	import java.util.Set;
	import java.util.stream.Collectors;

	import org.apache.helix.HelixConstants;
	import org.apache.helix.controller.dataproviders.ResourceControllerDataProvider;
	import org.apache.helix.controller.rebalancer.waged.WagedRebalancer;
	import org.apache.helix.model.CurrentState;
	import org.apache.helix.model.IdealState;
	import org.apache.helix.model.InstanceConfig;
	import org.apache.helix.model.LiveInstance;
	import org.apache.helix.model.Partition;
	import org.apache.helix.model.Resource;
	import org.apache.helix.model.ResourceAssignment;
	import org.mockito.stubbing.Answer;
	import org.testng.Assert;
	import org.testng.annotations.BeforeClass;
	import org.testng.annotations.Test;

	import static org.mockito.Matchers.anyString;
	import static org.mockito.Mockito.when;

	public class TestClusterModelProvider extends AbstractTestClusterModel {
	Set<String> _instances;

	@BeforeClass
	public void initialize() {
	super.initialize();
	_instances = new HashSet<>();
	_instances.add(_testInstanceId);
	}

	@Override
	protected ResourceControllerDataProvider setupClusterDataCache() throws IOException {
	ResourceControllerDataProvider testCache = super.setupClusterDataCache();

	// Set up mock idealstate
	Map<String, IdealState> isMap = new HashMap<>();
	for (String resource : _resourceNames) {
	IdealState is = new IdealState(resource);
	is.setNumPartitions(_partitionNames.size());
	is.setRebalanceMode(IdealState.RebalanceMode.FULL_AUTO);
	is.setStateModelDefRef("MasterSlave");
	is.setReplicas("3");
	is.setRebalancerClassName(WagedRebalancer.class.getName());
	_partitionNames.stream()
	.forEach(partition -> is.setPreferenceList(partition, Collections.emptyList()));
	isMap.put(resource, is);
	}
	when(testCache.getIdealState(anyString())).thenAnswer(
	(Answer<IdealState>) invocationOnMock -> isMap.get(invocationOnMock.getArguments()[0]));

	// Set up 2 more instances
	for (int i = 1; i < 3; i++) {
	String instanceName = _testInstanceId + i;
	_instances.add(instanceName);
	// 1. Set up the default instance information with capacity configuration.
	InstanceConfig testInstanceConfig = createMockInstanceConfig(instanceName);
	Map<String, InstanceConfig> instanceConfigMap = testCache.getInstanceConfigMap();
	instanceConfigMap.put(instanceName, testInstanceConfig);
	when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
	// 2. Mock the live instance node for the default instance.
	LiveInstance testLiveInstance = createMockLiveInstance(instanceName);
	Map<String, LiveInstance> liveInstanceMap = testCache.getLiveInstances();
	liveInstanceMap.put(instanceName, testLiveInstance);
	when(testCache.getLiveInstances()).thenReturn(liveInstanceMap);
	}

	return testCache;
	}

	@Test
	public void testGenerateClusterModel() throws IOException {
	ResourceControllerDataProvider testCache = setupClusterDataCache();
	// 1. test generating a cluster model with empty assignment
	ClusterModel clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, Collections.emptyMap(), Collections.emptyMap());
	// There should be no existing assignment.
	Assert.assertFalse(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.anyMatch(resourceMap -> !resourceMap.isEmpty()));
	Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
	.anyMatch(node -> node.getAssignedReplicaCount() != 0));
	// Have all 3 instances
	Assert.assertEquals(
	clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
	.collect(Collectors.toSet()), _instances);
	// Shall have 2 resources and 4 replicas, since all nodes are in the same fault zone.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 4));

	// Adjust instance fault zone, so they have different fault zones.
	testCache.getInstanceConfigMap().values().stream()
	.forEach(config -> config.setZoneId(config.getInstanceName()));
	clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, Collections.emptyMap(), Collections.emptyMap());
	// Shall have 2 resources and 12 replicas after fault zone adjusted.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 12));

	// 2. test with only one active node
	clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	Collections.singleton(_testInstanceId), Collections.emptyMap(), Collections.emptyMap());
	// Have only one instance
	Assert.assertEquals(
	clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
	.collect(Collectors.toSet()), Collections.singleton(_testInstanceId));
	// Shall have 4 assignable replicas because there is only one valid node.
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 4));

	// 3. test with no active instance
	clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	Collections.emptySet(), Collections.emptyMap(), Collections.emptyMap());
	// Have only one instance
	Assert.assertEquals(clusterModel.getAssignableNodes().size(), 0);
	// Shall have 0 assignable replicas because there are 0 valid nodes.
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.isEmpty()));

	// 4. test with baseline assignment
	// Mock a baseline assignment based on the current states.
	Map<String, ResourceAssignment> baselineAssignment = new HashMap<>();
	for (String resource : _resourceNames) {
	// <partition, <instance, state>>
	Map<String, Map<String, String>> assignmentMap = new HashMap<>();
	CurrentState cs = testCache.getCurrentState(_testInstanceId, _sessionId).get(resource);
	if (cs != null) {
	for (Map.Entry<String, String> stateEntry : cs.getPartitionStateMap().entrySet()) {
	assignmentMap.computeIfAbsent(stateEntry.getKey(), k -> new HashMap<>())
	.put(_testInstanceId, stateEntry.getValue());
	}
	ResourceAssignment assignment = new ResourceAssignment(resource);
	assignmentMap.keySet().stream().forEach(partition -> assignment
	.addReplicaMap(new Partition(partition), assignmentMap.get(partition)));
	baselineAssignment.put(resource, assignment);
	}
	}

	// Generate a cluster model based on the best possible assignment
	clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, Collections.emptyMap(), baselineAssignment);
	// There should be 4 existing assignments in total (each resource has 2) in the specified instance
	Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.allMatch(resourceMap -> resourceMap.values().stream()
	.allMatch(partitionSet -> partitionSet.size() == 2)));
	Assert.assertEquals(
	clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 4);
	// Since each resource has 2 replicas assigned, the assignable replica count should be 10.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 10));

	// 5. test with best possible assignment but cluster topology is changed
	clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances,
	Collections.singletonMap(HelixConstants.ChangeType.CLUSTER_CONFIG, Collections.emptySet()),
	baselineAssignment);
	// There should be no existing assignment since the topology change invalidates all existing assignment
	Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.allMatch(resourceMap -> resourceMap.isEmpty()));
	Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
	.anyMatch(node -> node.getAssignedReplicaCount() != 0));
	// Shall have 2 resources and 12 replicas
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 12));

	// 6. test with best possible assignment and one resource config change
	// Generate a cluster model based on the same best possible assignment, but resource1 config is changed
	String changedResourceName = _resourceNames.get(0);
	clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, Collections.singletonMap(HelixConstants.ChangeType.RESOURCE_CONFIG,
	Collections.singleton(changedResourceName)), baselineAssignment);
	// There should be no existing assignment for all the resource except for resource2
	Assert.assertEquals(clusterModel.getContext().getAssignmentForFaultZoneMap().size(), 1);
	Map<String, Set<String>> resourceAssignmentMap =
	clusterModel.getContext().getAssignmentForFaultZoneMap().get(_testInstanceId);
	// Should be only resource2 in the map
	Assert.assertEquals(resourceAssignmentMap.size(), 1);
	for (String resource : _resourceNames) {
	Assert
	.assertEquals(resourceAssignmentMap.getOrDefault(resource, Collections.emptySet()).size(),
	resource.equals(changedResourceName) ? 0 : 2);
	}
	// Only the first instance will have 2 assignment from resource2.
	for (String instance : _instances) {
	Assert.assertEquals(clusterModel.getAssignableNodes().get(instance).getAssignedReplicaCount(),
	instance.equals(_testInstanceId) ? 2 : 0);
	}
	// Shall have 2 resources and 12 replicas
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().keySet().size(), 2);
	for (String resource : _resourceNames) {
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().get(resource).size(),
	resource.equals(changedResourceName) ? 12 : 10);
	}

	// 7. test with best possible assignment but the instance becomes inactive
	// Generate a cluster model based on the best possible assignment, but the assigned node is disabled
	Set<String> limitedActiveInstances = new HashSet<>(_instances);
	limitedActiveInstances.remove(_testInstanceId);
	clusterModel = ClusterModelProvider.generateClusterModelForBaseline(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	limitedActiveInstances, Collections.emptyMap(), baselineAssignment);
	// There should be no existing assignment.
	Assert.assertFalse(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.anyMatch(resourceMap -> !resourceMap.isEmpty()));
	Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
	.anyMatch(node -> node.getAssignedReplicaCount() != 0));
	// Have only 2 instances
	Assert.assertEquals(
	clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
	.collect(Collectors.toSet()), limitedActiveInstances);
	// Since only 2 instances are active, we shall have 8 assignable replicas in each resource.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 8));
	}

	@Test (dependsOnMethods = "testGenerateClusterModel")
	public void testGenerateClusterModelForPartialRebalance() throws IOException {
	ResourceControllerDataProvider testCache = setupClusterDataCache();
	// 1. test generating a cluster model with empty assignment
	ClusterModel clusterModel = ClusterModelProvider
	.generateClusterModelForPartialRebalance(testCache, _resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, Collections.emptyMap(), Collections.emptyMap());
	// There should be no existing assignment.
	Assert.assertFalse(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.anyMatch(resourceMap -> !resourceMap.isEmpty()));
	Assert.assertFalse(clusterModel.getAssignableNodes().values().stream()
	.anyMatch(node -> node.getAssignedReplicaCount() != 0));
	// Have all 3 instances
	Assert.assertEquals(
	clusterModel.getAssignableNodes().values().stream().map(AssignableNode::getInstanceName)
	.collect(Collectors.toSet()), _instances);
	// Shall have 0 resources and 0 replicas since the baseline is empty. The partial rebalance
	// should not rebalance any replica.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 0);

	// Adjust instance fault zone, so they have different fault zones.
	testCache.getInstanceConfigMap().values().stream()
	.forEach(config -> config.setZoneId(config.getInstanceName()));

	// 2. test with a pair of identical best possible assignment and baseline assignment
	// Mock a best possible assignment based on the current states.
	Map<String, ResourceAssignment> bestPossibleAssignment = new HashMap<>();
	for (String resource : _resourceNames) {
	// <partition, <instance, state>>
	Map<String, Map<String, String>> assignmentMap = new HashMap<>();
	CurrentState cs = testCache.getCurrentState(_testInstanceId, _sessionId).get(resource);
	if (cs != null) {
	for (Map.Entry<String, String> stateEntry : cs.getPartitionStateMap().entrySet()) {
	assignmentMap.computeIfAbsent(stateEntry.getKey(), k -> new HashMap<>())
	.put(_testInstanceId, stateEntry.getValue());
	}
	ResourceAssignment assignment = new ResourceAssignment(resource);
	assignmentMap.keySet().stream().forEach(partition -> assignment
	.addReplicaMap(new Partition(partition), assignmentMap.get(partition)));
	bestPossibleAssignment.put(resource, assignment);
	}
	}
	Map<String, ResourceAssignment> baseline = new HashMap<>(bestPossibleAssignment);
	// Generate a cluster model for partial rebalance
	clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, baseline, bestPossibleAssignment);
	// There should be 4 existing assignments in total (each resource has 2) in the specified instance
	Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.allMatch(resourceMap -> resourceMap.values().stream()
	.allMatch(partitionSet -> partitionSet.size() == 2)));
	Assert.assertEquals(
	clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 4);
	// Since the best possible matches the baseline, no replica needs to be reassigned.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 0);

	// 3. test with inactive instance in the baseline and the best possible assignment
	Set<String> partialInstanceList = new HashSet<>(_instances);
	partialInstanceList.remove(_testInstanceId);
	clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	partialInstanceList, baseline, bestPossibleAssignment);
	// Have the other 2 active instances
	Assert.assertEquals(clusterModel.getAssignableNodes().size(), 2);
	// All the replicas in the existing assignment should be rebalanced.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 2);
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 2));
	// Shall have 0 assigned replicas
	Assert.assertTrue(clusterModel.getAssignableNodes().values().stream()
	.allMatch(assignableNode -> assignableNode.getAssignedReplicaCount() == 0));

	// 4. test with one resource that is only in the baseline
	String resourceInBaselineOnly = _resourceNames.get(0);
	Map<String, ResourceAssignment> partialBestPossibleAssignment =
	new HashMap<>(bestPossibleAssignment);
	partialBestPossibleAssignment.remove(resourceInBaselineOnly);
	// Generate a cluster mode with the adjusted best possible assignment
	clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, baseline, partialBestPossibleAssignment);
	// There should be 2 existing assignments in total in the specified instance
	Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.allMatch(resourceMap -> resourceMap.values().stream()
	.allMatch(partitionSet -> partitionSet.size() == 2)));
	// Only the replicas of one resource require rebalance
	Assert.assertEquals(
	clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 2);
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 1);
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().containsKey(resourceInBaselineOnly));
	Assert.assertTrue(clusterModel.getAssignableReplicaMap().values().stream()
	.allMatch(replicaSet -> replicaSet.size() == 2));

	// 5. test with one resource only in the best possible assignment
	String resourceInBestPossibleOnly = _resourceNames.get(1);
	Map<String, ResourceAssignment> partialBaseline = new HashMap<>(baseline);
	partialBaseline.remove(resourceInBestPossibleOnly);
	// Generate a cluster model with the adjusted baseline
	clusterModel = ClusterModelProvider.generateClusterModelForPartialRebalance(testCache,
	_resourceNames.stream()
	.collect(Collectors.toMap(resource -> resource, resource -> new Resource(resource))),
	_instances, partialBaseline, bestPossibleAssignment);
	// There should be 2 existing assignments in total and all of them require rebalance.
	Assert.assertTrue(clusterModel.getContext().getAssignmentForFaultZoneMap().values().stream()
	.allMatch(resourceMap -> resourceMap.values().stream()
	.allMatch(partitionSet -> partitionSet.size() == 2)));
	Assert.assertEquals(
	clusterModel.getAssignableNodes().get(_testInstanceId).getAssignedReplicaCount(), 2);
	// No need to rebalance the replicas that are not in the baseline yet.
	Assert.assertEquals(clusterModel.getAssignableReplicaMap().size(), 0);
	}
	}