Source:
public void testMoveMaintenancePriorityWithNoMaintenance() { mockDruidServer(druidServer1, "1", "normal", 30L, 100L, ImmutableMap.of(segment1.getIdentifier(), segment1, segment2.getIdentifier(), segment2)); mockDruidServer(druidServer2, "2", "normal", 0L, 100L, ImmutableMap.of(segment3.getIdentifier(), segment3, segment4.getIdentifier(), segment4)); mockDruidServer(druidServer3, "3", "normal", 0L, 100L, Collections.emptyMap()); EasyMock.replay(druidServer4); mockCoordinator(coordinator); BalancerStrategy strategy = EasyMock.createMock(BalancerStrategy.class); EasyMock.expect(strategy.pickSegmentToMove(anyObject())) .andReturn(new BalancerSegmentHolder(druidServer1, segment1)) .andReturn(new BalancerSegmentHolder(druidServer1, segment2)) .andReturn(new BalancerSegmentHolder(druidServer2, segment3)) .andReturn(new BalancerSegmentHolder(druidServer2, segment4)); EasyMock.expect(strategy.findNewSegmentHomeBalancer(anyObject(), anyObject())).andReturn(new ServerHolder(druidServer3, peon3)).anyTimes(); replay(strategy); DruidCoordinatorRuntimeParams params = defaultRuntimeParamsBuilder( ImmutableList.of(druidServer1, druidServer2, druidServer3), ImmutableList.of(peon1, peon2, peon3), ImmutableList.of(false, false, false) ) .withDynamicConfigs( CoordinatorDynamicConfig.builder() .withMaxSegmentsToMove( 3 ).withMaintenanceModeSegmentsPriority(9).build() ) .withBalancerStrategy(strategy) .build(); params = new DruidCoordinatorBalancerTester(coordinator).run(params); Assert.assertEquals(3L, params.getCoordinatorStats().getTieredStat("movedCount", "normal")); Assert.assertThat(peon3.getSegmentsToLoad(), is(equalTo(ImmutableSet.of(segment1, segment2, segment3)))); }

Target:
public void testMoveMaintenancePriorityWithNoMaintenance() { mockDruidServer(druidServer1, "1", "normal", 30L, 100L, Arrays.asList(segment1, segment2)); mockDruidServer(druidServer2, "2", "normal", 0L, 100L, Arrays.asList(segment3, segment4)); mockDruidServer(druidServer3, "3", "normal", 0L, 100L, Collections.emptyList()); EasyMock.replay(druidServer4); mockCoordinator(coordinator); BalancerStrategy strategy = EasyMock.createMock(BalancerStrategy.class); EasyMock.expect(strategy.pickSegmentToMove(anyObject())) .andReturn(new BalancerSegmentHolder(druidServer1, segment1)) .andReturn(new BalancerSegmentHolder(druidServer1, segment2)) .andReturn(new BalancerSegmentHolder(druidServer2, segment3)) .andReturn(new BalancerSegmentHolder(druidServer2, segment4)); EasyMock.expect(strategy.findNewSegmentHomeBalancer(anyObject(), anyObject())) .andReturn(new ServerHolder(druidServer3, peon3)) .anyTimes(); replay(strategy); DruidCoordinatorRuntimeParams params = defaultRuntimeParamsBuilder( ImmutableList.of(druidServer1, druidServer2, druidServer3), ImmutableList.of(peon1, peon2, peon3), ImmutableList.of(false, false, false) ) .withDynamicConfigs( CoordinatorDynamicConfig.builder().withMaxSegmentsToMove(3).withMaintenanceModeSegmentsPriority(9).build() ) .withBalancerStrategy(strategy) .build(); params = new DruidCoordinatorBalancerTester(coordinator).run(params); Assert.assertEquals(3L, params.getCoordinatorStats().getTieredStat("movedCount", "normal")); Assert.assertThat(peon3.getSegmentsToLoad(), is(equalTo(ImmutableSet.of(segment1, segment2, segment3)))); }