@Test(timeout=300000)
public void testBalancerDuringUpgrade() throws Exception {

  final int SEED = 0xFADED;
  Configuration conf = new HdfsConfiguration();
  conf.setLong(DFS_HEARTBEAT_INTERVAL_KEY, 1);
  conf.setInt(DFS_NAMENODE_HEARTBEAT_RECHECK_INTERVAL_KEY, 500);
  conf.setLong(DFSConfigKeys.DFS_NAMENODE_REDUNDANCY_INTERVAL_SECONDS_KEY, 1);
  conf.setLong(DFSConfigKeys.DFS_BALANCER_GETBLOCKS_MIN_BLOCK_SIZE_KEY, 1L);
  final int BLOCK_SIZE = 1024*1024;
  cluster = new MiniDFSCluster
      .Builder(conf)
      .numDataNodes(1)
      .storageCapacities(new long[] { BLOCK_SIZE * 10 })
      .storageTypes(new StorageType[] { DEFAULT })
      .storagesPerDatanode(1)
      .build();
  cluster.waitActive();
  // Create a file on the single DN
  final String METHOD_NAME = GenericTestUtils.getMethodName();
  final Path path1 = new Path("/" + METHOD_NAME + ".01.dat");
  DistributedFileSystem fs = cluster.getFileSystem();
  DFSTestUtil.createFile(fs, path1, BLOCK_SIZE, BLOCK_SIZE * 2, BLOCK_SIZE,
      (short) 1, SEED);
  // Add another DN with the same capacity, cluster is now unbalanced
  cluster.startDataNodes(conf, 1, true, null, null);
  cluster.triggerHeartbeats();
  Collection<URI> namenodes = DFSUtil.getInternalNsRpcUris(conf);
  // Run balancer
  final BalancerParameters p = BalancerParameters.DEFAULT;
  fs.setSafeMode(HdfsConstants.SafeModeAction.SAFEMODE_ENTER);
  fs.rollingUpgrade(HdfsConstants.RollingUpgradeAction.PREPARE);
  fs.setSafeMode(HdfsConstants.SafeModeAction.SAFEMODE_LEAVE);
  // Rolling upgrade should abort the balancer
  assertEquals(ExitStatus.UNFINALIZED_UPGRADE.getExitCode(),
      Balancer.run(namenodes, p, conf));
  // Should work with the -runDuringUpgrade flag.
  BalancerParameters.Builder b =
      new BalancerParameters.Builder();
  b.setRunDuringUpgrade(true);
  final BalancerParameters runDuringUpgrade = b.build();
  assertEquals(ExitStatus.SUCCESS.getExitCode(),
      Balancer.run(namenodes, runDuringUpgrade, conf));
  // Finalize the rolling upgrade
  fs.rollingUpgrade(HdfsConstants.RollingUpgradeAction.FINALIZE);
  // Should also work after finalization.
  assertEquals(ExitStatus.SUCCESS.getExitCode(),
      Balancer.run(namenodes, p, conf));
}