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 com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
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.util.DelayedRebalanceUtil;
import org.apache.helix.controller.rebalancer.waged.WagedRebalancer;
import org.apache.helix.model.ClusterConfig;
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.apache.helix.model.ResourceConfig;
import org.mockito.Mockito;
import org.mockito.stubbing.Answer;
import org.testng.Assert;
import org.testng.annotations.Test;

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

public class TestClusterModelProvider extends AbstractTestClusterModel {
  Map<String, ResourceConfig> _resourceConfigMap = new HashMap<>();

  @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.getAssignableInstanceConfigMap();
      instanceConfigMap.put(instanceName, testInstanceConfig);
      when(testCache.getAssignableInstanceConfigMap()).thenReturn(instanceConfigMap);
      when(testCache.getInstanceConfigMap()).thenReturn(instanceConfigMap);
      // 2. Mock the live instance node for the default instance.
      LiveInstance testLiveInstance = createMockLiveInstance(instanceName);
      Map<String, LiveInstance> liveInstanceMap = testCache.getAssignableLiveInstances();
      liveInstanceMap.put(instanceName, testLiveInstance);
      when(testCache.getAssignableLiveInstances()).thenReturn(liveInstanceMap);
    }

    return testCache;
  }

  @Test
  public void testFindToBeAssignedReplicasForMinActiveReplica() throws IOException {
    ResourceControllerDataProvider testCache = setupClusterDataCache();
    String instance1 = _testInstanceId;
    String offlineInstance = _testInstanceId + "1";
    String instance2 = _testInstanceId + "2";
    Map<String, LiveInstance> liveInstanceMap = new HashMap<>();
    liveInstanceMap.put(instance1, createMockLiveInstance(instance1));
    liveInstanceMap.put(instance2, createMockLiveInstance(instance2));
    Set<String> activeInstances = new HashSet<>();
    activeInstances.add(instance1);
    activeInstances.add(instance2);
    when(testCache.getAssignableLiveInstances()).thenReturn(liveInstanceMap);
    when(testCache.getEnabledLiveInstances()).thenReturn(activeInstances);

    // test 0, empty input
    Assert.assertEquals(
        DelayedRebalanceUtil.findToBeAssignedReplicasForMinActiveReplica(testCache, Collections.emptySet(),
            activeInstances, Collections.emptyMap(), new HashMap<>()),
        Collections.emptySet());

    // test 1, one partition under minActiveReplica
    Map<String, Map<String, Map<String, String>>> input = ImmutableMap.of(
        _resourceNames.get(0),
        ImmutableMap.of(
            _partitionNames.get(0), ImmutableMap.of("MASTER", instance1),
            _partitionNames.get(1), ImmutableMap.of("OFFLINE", offlineInstance)), // Partition2-MASTER
        _resourceNames.get(1),
        ImmutableMap.of(
            _partitionNames.get(2), ImmutableMap.of("MASTER", instance1),
            _partitionNames.get(3), ImmutableMap.of("SLAVE", instance1))
    );
    Map<String, Set<AssignableReplica>> replicaMap = new HashMap<>(); // to populate
    Map<String, ResourceAssignment> currentAssignment = new HashMap<>(); // to populate
    prepareData(input, replicaMap, currentAssignment, testCache, 1);

    Map<String, Set<AssignableReplica>> allocatedReplicas = new HashMap<>();
    Set<AssignableReplica> toBeAssignedReplicas =
        DelayedRebalanceUtil.findToBeAssignedReplicasForMinActiveReplica(testCache, replicaMap.keySet(), activeInstances,
            currentAssignment, allocatedReplicas);

    Assert.assertEquals(toBeAssignedReplicas.size(), 1);
    Assert.assertTrue(toBeAssignedReplicas.stream().map(AssignableReplica::toString).collect(Collectors.toSet())
        .contains("Resource1-Partition2-MASTER"));
    AssignableReplica replica = toBeAssignedReplicas.iterator().next();
    Assert.assertEquals(replica.getReplicaState(), "MASTER");
    Assert.assertEquals(replica.getPartitionName(), "Partition2");

    // test 2, no additional replica to be assigned
    testCache = setupClusterDataCache();
    when(testCache.getAssignableLiveInstances()).thenReturn(liveInstanceMap);
    when(testCache.getEnabledLiveInstances()).thenReturn(activeInstances);
    input = ImmutableMap.of(
        _resourceNames.get(0),
        ImmutableMap.of(
            _partitionNames.get(0), ImmutableMap.of("MASTER", instance1),
            _partitionNames.get(1), ImmutableMap.of("SLAVE", instance2)),
        _resourceNames.get(1),
        ImmutableMap.of(
            _partitionNames.get(2), ImmutableMap.of("MASTER", instance1),
            _partitionNames.get(3), ImmutableMap.of("SLAVE", instance2))
    );
    replicaMap = new HashMap<>(); // to populate
    currentAssignment = new HashMap<>(); // to populate
    prepareData(input, replicaMap, currentAssignment, testCache, 1);
    allocatedReplicas = new HashMap<>();
    toBeAssignedReplicas =
        DelayedRebalanceUtil.findToBeAssignedReplicasForMinActiveReplica(testCache, replicaMap.keySet(), activeInstances,
            currentAssignment, allocatedReplicas);
    Assert.assertTrue(toBeAssignedReplicas.isEmpty());
    Assert.assertEquals(allocatedReplicas.get(instance1).size(), 2);
    Assert.assertEquals(allocatedReplicas.get(instance2).size(), 2);

    // test 3, minActiveReplica==2, two partitions falling short
    testCache = setupClusterDataCache();
    when(testCache.getAssignableLiveInstances()).thenReturn(liveInstanceMap);
    when(testCache.getEnabledLiveInstances()).thenReturn(activeInstances);
    input = ImmutableMap.of(
        _resourceNames.get(0),
        ImmutableMap.of(
            _partitionNames.get(0), ImmutableMap.of("MASTER", instance1, "SLAVE", instance2),
            _partitionNames.get(1), ImmutableMap.of("MASTER", instance1, "OFFLINE", offlineInstance)), // Partition2-SLAVE
        _resourceNames.get(1),
        ImmutableMap.of(
            _partitionNames.get(2), ImmutableMap.of("MASTER", instance1, "SLAVE", instance2),
            _partitionNames.get(3), ImmutableMap.of("SLAVE", instance1, "OFFLINE", offlineInstance)) // Partition4-MASTER
    );
    replicaMap = new HashMap<>(); // to populate
    currentAssignment = new HashMap<>(); // to populate
    prepareData(input, replicaMap, currentAssignment, testCache, 2);
    allocatedReplicas = new HashMap<>();
    toBeAssignedReplicas =
        DelayedRebalanceUtil.findToBeAssignedReplicasForMinActiveReplica(testCache, replicaMap.keySet(), activeInstances,
            currentAssignment, allocatedReplicas);
    Assert.assertEquals(toBeAssignedReplicas.size(), 2);
    Assert.assertEquals(toBeAssignedReplicas.stream().map(AssignableReplica::toString).collect(Collectors.toSet()),
        ImmutableSet.of("Resource1-Partition2-SLAVE", "Resource2-Partition4-MASTER"));
    Assert.assertEquals(allocatedReplicas.get(instance1).size(), 4);
    Assert.assertEquals(allocatedReplicas.get(instance2).size(), 2);
  }

  @Test(dependsOnMethods = "testFindToBeAssignedReplicasForMinActiveReplica")
  public void testClusterModelForDelayedRebalanceOverwrite() throws IOException {
    ResourceControllerDataProvider testCache = setupClusterDataCache();
    String instance1 = _testInstanceId;
    String offlineInstance = _testInstanceId + "1";
    String instance2 = _testInstanceId + "2";
    Map<String, LiveInstance> liveInstanceMap = new HashMap<>();
    liveInstanceMap.put(instance1, createMockLiveInstance(instance1));
    liveInstanceMap.put(instance2, createMockLiveInstance(instance2));
    Set<String> activeInstances = new HashSet<>();
    activeInstances.add(instance1);
    activeInstances.add(instance2);
    when(testCache.getAssignableLiveInstances()).thenReturn(liveInstanceMap);
    when(testCache.getEnabledLiveInstances()).thenReturn(activeInstances);

    // test 1, one partition under minActiveReplica
    Map<String, Map<String, Map<String, String>>> input = ImmutableMap.of(
        _resourceNames.get(0),
        ImmutableMap.of(
            _partitionNames.get(0), ImmutableMap.of("MASTER", instance1),
            _partitionNames.get(1), ImmutableMap.of("OFFLINE", offlineInstance), // Partition2-MASTER
            _partitionNames.get(2), ImmutableMap.of("MASTER", instance2),
            _partitionNames.get(3), ImmutableMap.of("MASTER", instance2)),
        _resourceNames.get(1),
        ImmutableMap.of(
            _partitionNames.get(0), ImmutableMap.of("MASTER", instance2),
            _partitionNames.get(1), ImmutableMap.of("MASTER", instance2),
            _partitionNames.get(2), ImmutableMap.of("MASTER", instance1),
            _partitionNames.get(3), ImmutableMap.of("OFFLINE", offlineInstance)) // Partition4-MASTER
    );
    Map<String, Set<AssignableReplica>> replicaMap = new HashMap<>(); // to populate
    Map<String, ResourceAssignment> currentAssignment = new HashMap<>(); // to populate
    prepareData(input, replicaMap, currentAssignment, testCache, 1);

    Map<String, Resource> resourceMap = _resourceNames.stream().collect(Collectors.toMap(resource -> resource, Resource::new));
    ClusterModel clusterModel = ClusterModelProvider.generateClusterModelForDelayedRebalanceOverwrites(testCache,
        resourceMap, activeInstances, currentAssignment);
    Assert.assertEquals(clusterModel.getAssignableNodes().size(), 2);
    Assert.assertTrue(clusterModel.getAssignableNodes().containsKey(instance1));
    Assert.assertTrue(clusterModel.getAssignableNodes().containsKey(instance2));
    Assert.assertEquals(clusterModel.getAssignableNodes().get(instance1).getAssignedReplicas().size(), 2);
    Assert.assertEquals(clusterModel.getAssignableNodes().get(instance2).getAssignedReplicas().size(), 4);

    Assert.assertEquals(clusterModel.getAssignableReplicaMap().get("Resource1").size(), 1);
    Assert.assertEquals(clusterModel.getAssignableReplicaMap().get("Resource1").iterator().next().toString(),
        "Resource1-Partition2-MASTER");
    Assert.assertEquals(clusterModel.getAssignableReplicaMap().get("Resource2").size(), 1);
    Assert.assertEquals(clusterModel.getAssignableReplicaMap().get("Resource2").iterator().next().toString(),
        "Resource2-Partition4-MASTER");

    // test 2, minActiveReplica==2, three partitions falling short
    testCache = setupClusterDataCache();
    when(testCache.getAssignableLiveInstances()).thenReturn(liveInstanceMap);
    when(testCache.getEnabledLiveInstances()).thenReturn(activeInstances);
    input = ImmutableMap.of(
        _resourceNames.get(0),
        ImmutableMap.of(
            _partitionNames.get(0), ImmutableMap.of("MASTER", instance1, "SLAVE", instance2),
            _partitionNames.get(1), ImmutableMap.of("MASTER", instance1, "OFFLINE", offlineInstance), // Partition2-SLAVE
            _partitionNames.get(2), ImmutableMap.of("OFFLINE", offlineInstance, "SLAVE", instance2), // Partition3-MASTER
            _partitionNames.get(3), ImmutableMap.of("MASTER", instance1, "SLAVE", instance2)),
        _resourceNames.get(1),
        ImmutableMap.of(
            _partitionNames.get(0), ImmutableMap.of("MASTER", instance1, "SLAVE", instance2),
            _partitionNames.get(1), ImmutableMap.of("MASTER", instance1, "SLAVE", instance2),
            _partitionNames.get(2), ImmutableMap.of("MASTER", instance1, "SLAVE", instance2),
            _partitionNames.get(3), ImmutableMap.of("OFFLINE", offlineInstance, "ERROR", instance2)) // Partition4-MASTER
    );
    replicaMap = new HashMap<>(); // to populate
    currentAssignment = new HashMap<>(); // to populate
    prepareData(input, replicaMap, currentAssignment, testCache, 2);
    clusterModel = ClusterModelProvider.generateClusterModelForDelayedRebalanceOverwrites(testCache,
        resourceMap, activeInstances, currentAssignment);
    Assert.assertEquals(clusterModel.getAssignableNodes().size(), 2);
    Assert.assertTrue(clusterModel.getAssignableNodes().containsKey(instance1));
    Assert.assertTrue(clusterModel.getAssignableNodes().containsKey(instance2));
    Assert.assertEquals(clusterModel.getAssignableNodes().get(instance1).getAssignedReplicas().size(), 6);
    Assert.assertEquals(clusterModel.getAssignableNodes().get(instance2).getAssignedReplicas().size(), 7);

    Set<String> replicaSet = clusterModel.getAssignableReplicaMap().get(_resourceNames.get(0))
        .stream()
        .map(AssignableReplica::toString)
        .collect(Collectors.toSet());
    Assert.assertEquals(replicaSet.size(), 2);
    Assert.assertTrue(replicaSet.contains("Resource1-Partition2-SLAVE"));
    Assert.assertTrue(replicaSet.contains("Resource1-Partition3-MASTER"));
    replicaSet = clusterModel.getAssignableReplicaMap().get(_resourceNames.get(1))
        .stream()
        .map(AssignableReplica::toString)
        .collect(Collectors.toSet());
    Assert.assertEquals(replicaSet.size(), 1);
    Assert.assertTrue(replicaSet.contains("Resource2-Partition4-MASTER"));
  }

  /**
   * Prepare mock objects with given input. This methods prepare replicaMap and populate testCache with currentState.
   *
   * @param input <resource, <partition, <state, instance> > >
   * @param replicaMap The data map to prepare, a set of AssignableReplica by resource name.
   * @param currentAssignment The data map to prepare, resourceAssignment by resource name
   * @param testCache The mock object to prepare
   */
  private void prepareData(Map<String, Map<String, Map<String, String>>> input,
      Map<String, Set<AssignableReplica>> replicaMap,
      Map<String, ResourceAssignment> currentAssignment,
      ResourceControllerDataProvider testCache,
      int minActiveReplica) {

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


    // <instance, <resource, CurrentState>>
    Map<String, Map<String, CurrentState>> currentStateByInstanceByResource = new HashMap<>();
    Map<String, Map<String, Map<String, String>>> stateByInstanceByResourceByPartition = new HashMap<>();

    for (String resource : input.keySet()) {
      Set<AssignableReplica> replicas = new HashSet<>();
      replicaMap.put(resource, replicas);
      ResourceConfig resourceConfig = _resourceConfigMap.get(resource);
      for (String partition : input.get(resource).keySet()) {
        input.get(resource).get(partition).forEach(
            (state, instance) -> {
              stateByInstanceByResourceByPartition
                  .computeIfAbsent(instance, k -> new HashMap<>())
                  .computeIfAbsent(resource, k -> new HashMap<>())
                  .put(partition, state);
              replicas.add(new MockAssignableReplica(resourceConfig, partition, state));
            });
      }
    }
    for (String instance : stateByInstanceByResourceByPartition.keySet()) {
      for (String resource : stateByInstanceByResourceByPartition.get(instance).keySet()) {
        Map<String, String> partitionState = stateByInstanceByResourceByPartition.get(instance).get(resource);
        CurrentState testCurrentStateResource = mockCurrentStateResource(partitionState);
        currentStateByInstanceByResource.computeIfAbsent(instance, k -> new HashMap<>()).put(resource, testCurrentStateResource);
      }
    }

    for (String instance : currentStateByInstanceByResource.keySet()) {
      when(testCache.getCurrentState(instance, _sessionId)).thenReturn(currentStateByInstanceByResource.get(instance));
      when(testCache.getCurrentState(instance, _sessionId, false))
          .thenReturn(currentStateByInstanceByResource.get(instance));
    }

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

  private CurrentState mockCurrentStateResource(Map<String, String> partitionState) {
    CurrentState testCurrentStateResource = Mockito.mock(CurrentState.class);
    when(testCurrentStateResource.getResourceName()).thenReturn(_resourceNames.get(0));
    when(testCurrentStateResource.getPartitionStateMap()).thenReturn(partitionState);
    when(testCurrentStateResource.getStateModelDefRef()).thenReturn("MasterSlave");
    when(testCurrentStateResource.getSessionId()).thenReturn(_sessionId);
    for (Map.Entry<String, String> entry : partitionState.entrySet()) {
      when(testCurrentStateResource.getState(entry.getKey())).thenReturn(entry.getValue());
    }
    return testCurrentStateResource;
  }

  @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.getAssignableInstanceConfigMap().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.getAssignableInstanceConfigMap().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);
  }

  static class MockAssignableReplica extends AssignableReplica {
    MockAssignableReplica(ResourceConfig resourceConfig, String partition, String replicaState) {
      super(new ClusterConfig("testCluster"), resourceConfig, partition, replicaState, 1);
    }
  }
}
