BandwidthPartitioner extends StatelessPartitioner<T> { @Override public Collection<Partition<T>> definePartitions(Collection<Partition<T>> partitions, PartitioningContext context) { long currentBandwidth = partitions.iterator().next().getPartitionedInstance().getBandwidthManager().getBandwidth() * partitions.size(); Collection<Partition<T>> newpartitions = super.definePartitions(partitions, context); return updateBandwidth(newpartitions, currentBandwidth); } BandwidthPartitioner();  BandwidthPartitioner(String value);  BandwidthPartitioner(int partitionCount); @Override Collection<Partition<T>> definePartitions(Collection<Partition<T>> partitions, PartitioningContext context); Collection<Partition<T>> updateBandwidth(Collection<Partition<T>> newpartitions, long currentBandwidth);  }
@Test public void testBandwidthOnPartitions() { when(partitionContextMock.getParallelPartitionCount()).thenReturn(0); Collection<Partition<BandwidthLimitingOperator>> partitions = Lists.newArrayList(); DefaultPartition<BandwidthLimitingOperator> defaultPartition = new DefaultPartition<BandwidthLimitingOperator>(operatorMock); partitions.add(defaultPartition); underTest.definePartitions(partitions, partitionContextMock); verify(bandwidthManagerMock).setBandwidth(10L); }
@Test public void testBandwidthOnIncresedPartitions() { when(partitionContextMock.getParallelPartitionCount()).thenReturn(5); Collection<Partition<BandwidthLimitingOperator>> partitions = Lists.newArrayList(); DefaultPartition<BandwidthLimitingOperator> defaultPartition = new DefaultPartition<BandwidthLimitingOperator>(operatorMock); partitions.add(defaultPartition); underTest.definePartitions(partitions, partitionContextMock); verify(bandwidthManagerMock, times(5)).setBandwidth(2L); }
@Test public void testBandwidthOnReducedPartitions() { when(partitionContextMock.getParallelPartitionCount()).thenReturn(2); when(bandwidthManagerMock.getBandwidth()).thenReturn(2L); Collection<Partition<BandwidthLimitingOperator>> partitions = Lists.newArrayList(); for (int i = 5; i-- > 0;) { partitions.add(new DefaultPartition<BandwidthLimitingOperator>(operatorMock)); } underTest.definePartitions(partitions, partitionContextMock); verify(bandwidthManagerMock, times(2)).setBandwidth(5L); }