@Test
  public void optimalLeaderDistributionTest() {
    // Prepare Data
    List<TConsensusGroupId> regionGroupIds = new ArrayList<>();
    for (int i = 0; i < 3; i++) {
      regionGroupIds.add(new TConsensusGroupId(TConsensusGroupType.DataRegion, i));
    }
    List<TDataNodeLocation> dataNodeLocations = new ArrayList<>();
    for (int i = 0; i < 4; i++) {
      dataNodeLocations.add(new TDataNodeLocation().setDataNodeId(i));
    }
    // DataNode-0: [0, 1, 2], DataNode-1: [0, 1]
    // DataNode-2: [0, 2]   , DataNode-3: [1, 2]
    // The result will be unbalanced if select DataNode-2 as leader for RegionGroup-0
    // and select DataNode-3 as leader for RegionGroup-1
    List<TRegionReplicaSet> regionReplicaSets = new ArrayList<>();
    Map<TConsensusGroupId, Map<Integer, RegionStatistics>> regionStatisticsMap = new TreeMap<>();
    regionReplicaSets.add(
        new TRegionReplicaSet(
            regionGroupIds.get(0),
            Arrays.asList(
                dataNodeLocations.get(0), dataNodeLocations.get(1), dataNodeLocations.get(2))));
    Map<Integer, RegionStatistics> region0 = new TreeMap<>();
    region0.put(0, new RegionStatistics(RegionStatus.Unknown));
    region0.put(1, new RegionStatistics(RegionStatus.Running));
    region0.put(2, new RegionStatistics(RegionStatus.Running));
    regionStatisticsMap.put(regionGroupIds.get(0), region0);
    regionReplicaSets.add(
        new TRegionReplicaSet(
            regionGroupIds.get(1),
            Arrays.asList(
                dataNodeLocations.get(0), dataNodeLocations.get(1), dataNodeLocations.get(3))));
    Map<Integer, RegionStatistics> region1 = new TreeMap<>();
    region1.put(0, new RegionStatistics(RegionStatus.Unknown));
    region1.put(1, new RegionStatistics(RegionStatus.Running));
    region1.put(3, new RegionStatistics(RegionStatus.Running));
    regionStatisticsMap.put(regionGroupIds.get(1), region1);
    regionReplicaSets.add(
        new TRegionReplicaSet(
            regionGroupIds.get(2),
            Arrays.asList(
                dataNodeLocations.get(0), dataNodeLocations.get(2), dataNodeLocations.get(3))));
    Map<Integer, RegionStatistics> region2 = new TreeMap<>();
    region2.put(0, new RegionStatistics(RegionStatus.Unknown));
    region2.put(2, new RegionStatistics(RegionStatus.Running));
    region2.put(3, new RegionStatistics(RegionStatus.Running));
    regionStatisticsMap.put(regionGroupIds.get(2), region2);

    // Prepare input parameters
    Map<String, List<TConsensusGroupId>> databaseRegionGroupMap = new TreeMap<>();
    databaseRegionGroupMap.put(DATABASE, regionGroupIds);
    Map<TConsensusGroupId, Set<Integer>> regionReplicaSetMap = new TreeMap<>();
    regionReplicaSets.forEach(
        regionReplicaSet ->
            regionReplicaSetMap.put(
                regionReplicaSet.getRegionId(),
                regionReplicaSet.getDataNodeLocations().stream()
                    .map(TDataNodeLocation::getDataNodeId)
                    .collect(Collectors.toSet())));
    Map<TConsensusGroupId, Integer> regionLeaderMap = new TreeMap<>();
    regionReplicaSets.forEach(
        regionReplicaSet -> regionLeaderMap.put(regionReplicaSet.getRegionId(), 0));
    Map<Integer, NodeStatistics> dataNodeStatisticsMap = new TreeMap<>();
    dataNodeStatisticsMap.put(0, new NodeStatistics(NodeStatus.Unknown));
    dataNodeStatisticsMap.put(1, new NodeStatistics(NodeStatus.Running));
    dataNodeStatisticsMap.put(2, new NodeStatistics(NodeStatus.Running));
    dataNodeStatisticsMap.put(3, new NodeStatistics(NodeStatus.Running));

    // Do balancing
    Map<TConsensusGroupId, Integer> leaderDistribution =
        BALANCER.generateOptimalLeaderDistribution(
            databaseRegionGroupMap,
            regionReplicaSetMap,
            regionLeaderMap,
            dataNodeStatisticsMap,
            regionStatisticsMap);
    // All RegionGroup got a leader
    Assert.assertEquals(3, leaderDistribution.size());
    // Each DataNode has exactly one leader
    Assert.assertEquals(3, new HashSet<>(leaderDistribution.values()).size());
    // MaxFlow is 3
    Assert.assertEquals(3, BALANCER.getMaximumFlow());
    // MinimumCost is 3(switch leader cost) + 3(load cost, rNode -> sDNode)
    // + 3(load cost, sDNode -> tDNode)
    Assert.assertEquals(3 + 3 + 3, BALANCER.getMinimumCost());
  }
This test shows a simple case that greedy algorithm might fail
iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/balancer/router/leader/CFDLeaderBalancerTest.java