Source:
public void testDropDuringMaintenance() { final LoadQueuePeon mockPeon = createEmptyPeon(); mockPeon.dropSegment(EasyMock.anyObject(), EasyMock.anyObject()); EasyMock.expectLastCall().times(2); EasyMock.expect(mockBalancerStrategy.pickServersToDrop(EasyMock.anyObject(), EasyMock.anyObject())) .andDelegateTo(balancerStrategy) .times(2); EasyMock.replay(throttler, mockPeon, mockBalancerStrategy); LoadRule rule = createLoadRule(ImmutableMap.of( "tier1", 0 )); final DataSegment segment1 = createDataSegment("foo1"); final DataSegment segment2 = createDataSegment("foo2"); DruidServer server1 = createServer("tier1"); server1.addDataSegment(segment1); DruidServer server2 = createServer("tier1"); server2.addDataSegment(segment2); DruidCluster druidCluster = new DruidCluster( null, ImmutableMap.of( "tier1", Stream.of( new ServerHolder( server1.toImmutableDruidServer(), mockPeon, true ), new ServerHolder( server2.toImmutableDruidServer(), mockPeon, false ) ).collect(Collectors.toCollection(() -> new TreeSet<>(Collections.reverseOrder()))) ) ); DruidCoordinatorRuntimeParams params = DruidCoordinatorRuntimeParams.newBuilder() .withDruidCluster(druidCluster) .withSegmentReplicantLookup( SegmentReplicantLookup.make(druidCluster)) .withReplicationManager(throttler) .withBalancerStrategy(mockBalancerStrategy) .withBalancerReferenceTimestamp(DateTimes.of( "2013-01-01")) .withAvailableSegments(Arrays.asList( segment1, segment2 )).build(); CoordinatorStats stats = rule.run( null, params, segment1 ); Assert.assertEquals(1L, stats.getTieredStat("droppedCount", "tier1")); stats = rule.run( null, params, segment2 ); Assert.assertEquals(1L, stats.getTieredStat("droppedCount", "tier1")); EasyMock.verify(throttler, mockPeon); }

Target:
public void testDropDuringMaintenance() { final LoadQueuePeon mockPeon = createEmptyPeon(); mockPeon.dropSegment(EasyMock.anyObject(), EasyMock.anyObject()); EasyMock.expectLastCall().times(2); EasyMock.expect(mockBalancerStrategy.pickServersToDrop(EasyMock.anyObject(), EasyMock.anyObject())) .andDelegateTo(balancerStrategy) .times(4); EasyMock.replay(throttler, mockPeon, mockBalancerStrategy); LoadRule rule = createLoadRule(ImmutableMap.of("tier1", 0)); final DataSegment segment1 = createDataSegment("foo1"); final DataSegment segment2 = createDataSegment("foo2"); DruidServer server1 = createServer("tier1"); server1.addDataSegment(segment1); DruidServer server2 = createServer("tier1"); server2.addDataSegment(segment2); DruidCluster druidCluster = new DruidCluster( null, ImmutableMap.of( "tier1", Arrays.asList( new ServerHolder(server1.toImmutableDruidServer(), mockPeon, true), new ServerHolder(server2.toImmutableDruidServer(), mockPeon, false) ) ) ); DruidCoordinatorRuntimeParams params = DruidCoordinatorRuntimeParams .newBuilder() .withDruidCluster(druidCluster) .withSegmentReplicantLookup(SegmentReplicantLookup.make(druidCluster)) .withReplicationManager(throttler) .withBalancerStrategy(mockBalancerStrategy) .withBalancerReferenceTimestamp(DateTimes.of("2013-01-01")) .withAvailableSegments(segment1, segment2) .build(); CoordinatorStats stats = rule.run( null, params, segment1 ); Assert.assertEquals(1L, stats.getTieredStat("droppedCount", "tier1")); stats = rule.run( null, params, segment2 ); Assert.assertEquals(1L, stats.getTieredStat("droppedCount", "tier1")); EasyMock.verify(throttler, mockPeon); }