@Test
  public void migrateTest() {
    // All DataNodes are in Running status
    Map<Integer, NodeStatistics> nodeStatisticsMap = new TreeMap<>();
    nodeStatisticsMap.put(0, new NodeStatistics(NodeStatus.Running));
    nodeStatisticsMap.put(1, new NodeStatistics(NodeStatus.Running));
    // Prepare RegionGroups
    Map<String, List<TConsensusGroupId>> databaseRegionGroupMap = new TreeMap<>();
    Map<TConsensusGroupId, Set<Integer>> regionReplicaSetMap = new TreeMap<>();
    Map<TConsensusGroupId, Integer> regionLeaderMap = new TreeMap<>();
    Map<TConsensusGroupId, Map<Integer, RegionStatistics>> regionStatisticsMap = new TreeMap<>();
    for (int i = 0; i < 5; i++) {
      TConsensusGroupId regionGroupId = new TConsensusGroupId(TConsensusGroupType.DataRegion, i);
      databaseRegionGroupMap
          .computeIfAbsent(DATABASE, empty -> new ArrayList<>())
          .add(regionGroupId);
      TRegionReplicaSet regionReplicaSet =
          new TRegionReplicaSet(
              regionGroupId,
              Arrays.asList(
                  new TDataNodeLocation().setDataNodeId(0),
                  new TDataNodeLocation().setDataNodeId(1)));
      regionReplicaSetMap.put(
          regionGroupId,
          regionReplicaSet.getDataNodeLocations().stream()
              .map(TDataNodeLocation::getDataNodeId)
              .collect(Collectors.toSet()));
      regionLeaderMap.put(regionGroupId, 0);
      // Assuming all Regions are migrating from DataNode-1 to DataNode-2
      Map<Integer, RegionStatistics> regionStatistics = new TreeMap<>();
      regionStatistics.put(0, new RegionStatistics(RegionStatus.Removing));
      regionStatistics.put(1, new RegionStatistics(RegionStatus.Running));
      regionStatisticsMap.put(regionGroupId, regionStatistics);
    }

    // Do balancing
    Map<TConsensusGroupId, Integer> leaderDistribution =
        BALANCER.generateOptimalLeaderDistribution(
            databaseRegionGroupMap,
            regionReplicaSetMap,
            regionLeaderMap,
            nodeStatisticsMap,
            regionStatisticsMap);
    for (int i = 0; i < 5; i++) {
      // All RegionGroups' leader should be DataNode-1
      Assert.assertEquals(
          1,
          leaderDistribution
              .get(new TConsensusGroupId(TConsensusGroupType.DataRegion, i))
              .intValue());
    }
  }
The leader will remain the same if all DataNodes are disabled
iotdb-core/confignode/src/test/java/org/apache/iotdb/confignode/manager/load/balancer/router/leader/CFDLeaderBalancerTest.java