Correct getRandomBalancerSegmentHolderTest (#10569)

This commit is contained in:
Lucas Capistrant 2020-11-12 11:51:05 -06:00 committed by GitHub
parent 6563599de4
commit c3cad461bc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -137,49 +137,31 @@ public class ReservoirSegmentSamplerTest
@Test
public void getRandomBalancerSegmentHolderTest()
{
EasyMock.expect(druidServer1.getType()).andReturn(ServerType.HISTORICAL).atLeastOnce();
EasyMock.expect(druidServer1.getName()).andReturn("1").atLeastOnce();
EasyMock.expect(druidServer1.getCurrSize()).andReturn(30L).atLeastOnce();
EasyMock.expect(druidServer1.getMaxSize()).andReturn(100L).atLeastOnce();
int iterations = 5000;
EasyMock.expect(druidServer1.getType()).andReturn(ServerType.HISTORICAL).times(iterations);
ImmutableDruidServerTests.expectSegments(druidServer1, segments1);
EasyMock.expect(druidServer1.getSegment(EasyMock.anyObject())).andReturn(null).anyTimes();
EasyMock.replay(druidServer1);
EasyMock.expect(druidServer2.getType()).andReturn(ServerType.HISTORICAL).atLeastOnce();
EasyMock.expect(druidServer2.getName()).andReturn("2").atLeastOnce();
EasyMock.expect(druidServer2.getTier()).andReturn("normal").anyTimes();
EasyMock.expect(druidServer2.getCurrSize()).andReturn(30L).atLeastOnce();
EasyMock.expect(druidServer2.getMaxSize()).andReturn(100L).atLeastOnce();
EasyMock.expect(druidServer2.getType()).andReturn(ServerType.HISTORICAL).times(iterations);
ImmutableDruidServerTests.expectSegments(druidServer2, segments2);
EasyMock.expect(druidServer2.getSegment(EasyMock.anyObject())).andReturn(null).anyTimes();
EasyMock.replay(druidServer2);
EasyMock.expect(druidServer3.getType()).andReturn(ServerType.HISTORICAL).atLeastOnce();
EasyMock.expect(druidServer3.getName()).andReturn("3").atLeastOnce();
EasyMock.expect(druidServer3.getTier()).andReturn("normal").anyTimes();
EasyMock.expect(druidServer3.getCurrSize()).andReturn(30L).atLeastOnce();
EasyMock.expect(druidServer3.getMaxSize()).andReturn(100L).atLeastOnce();
EasyMock.expect(druidServer3.getType()).andReturn(ServerType.HISTORICAL).times(iterations);
ImmutableDruidServerTests.expectSegments(druidServer3, segments3);
EasyMock.expect(druidServer3.getSegment(EasyMock.anyObject())).andReturn(null).anyTimes();
EasyMock.replay(druidServer3);
EasyMock.expect(druidServer4.getType()).andReturn(ServerType.HISTORICAL).atLeastOnce();
EasyMock.expect(druidServer4.getName()).andReturn("4").atLeastOnce();
EasyMock.expect(druidServer4.getTier()).andReturn("normal").anyTimes();
EasyMock.expect(druidServer4.getCurrSize()).andReturn(30L).atLeastOnce();
EasyMock.expect(druidServer4.getMaxSize()).andReturn(100L).atLeastOnce();
EasyMock.expect(druidServer4.getType()).andReturn(ServerType.HISTORICAL).times(iterations);
ImmutableDruidServerTests.expectSegments(druidServer4, segments4);
EasyMock.expect(druidServer4.getSegment(EasyMock.anyObject())).andReturn(null).anyTimes();
EasyMock.replay(druidServer4);
// Have to use anyTimes() because the number of times a segment on a given server is chosen is indetermistic.
EasyMock.expect(holder1.getServer()).andReturn(druidServer1).anyTimes();
EasyMock.replay(holder1);
EasyMock.expect(holder2.getServer()).andReturn(druidServer2).anyTimes();
EasyMock.replay(holder2);
EasyMock.expect(holder3.getServer()).andReturn(druidServer3).anyTimes();
EasyMock.replay(holder3);
EasyMock.expect(holder4.getServer()).andReturn(druidServer4).anyTimes();
EasyMock.replay(holder4);
@ -190,12 +172,19 @@ public class ReservoirSegmentSamplerTest
holderList.add(holder4);
Map<DataSegment, Integer> segmentCountMap = new HashMap<>();
for (int i = 0; i < 5000; i++) {
segmentCountMap.put(ReservoirSegmentSampler.getRandomBalancerSegmentHolder(holderList, Collections.emptySet()).getSegment(), 1);
for (int i = 0; i < iterations; i++) {
// due to the pseudo-randomness of this method, we may not select a segment every single time no matter what.
BalancerSegmentHolder balancerSegmentHolder = ReservoirSegmentSampler.getRandomBalancerSegmentHolder(holderList, Collections.emptySet());
if (balancerSegmentHolder != null) {
segmentCountMap.put(balancerSegmentHolder.getSegment(), 1);
}
}
for (DataSegment segment : segments) {
Assert.assertEquals(segmentCountMap.get(segment), new Integer(1));
}
EasyMock.verify(druidServer1, druidServer2, druidServer3, druidServer4);
EasyMock.verify(holder1, holder2, holder3, holder4);
}
}