EmployeeJobPartitioner implements Partitioner { @Override public Map<String, ExecutionContext> partition(int gridSize) { int partitionCount = gridSize * PARTITIONS_PER_NODE; List<Long> employeeIds = taxCalculationRepository.getUnprocessedEmployeeIds(year, month, stepExecution.getJobExecutionId()); int size = employeeIds.size(); int targetSize = size / partitionCount + 1; Map<String, ExecutionContext> result = new HashMap<>(); for (int index = 0, partitionNumber = 0; index < size; index += targetSize, partitionNumber++) { ExecutionContext value = new ExecutionContext(); value.putLong("minValue", employeeIds.get(index)); value.putLong("maxValue", employeeIds.get(Math.min(index + targetSize - 1, size - 1))); value.putLong("partition", partitionNumber); result.put("partition" + partitionNumber, value); } return result; }  @Override Map<String, ExecutionContext> partition(int gridSize);  }
@Test public void given31Employees_whenPartitioningFor1NodeWith5PartitionsPerNode_thenDistributeCorrectly() throws Exception { List<Long> employeeIds = new ArrayList<>(); for (long i = 1; i <= 31; i++) { employeeIds.add(i); } when(taxCalculationRepository.getUnprocessedEmployeeIds(anyLong(), anyLong(), anyLong())).thenReturn(employeeIds); Map<String, ExecutionContext> partitions = employeeJobPartitioner.partition(1); assertThat(partitions).hasSize(5); assertThat(partitions.get("partition0").getLong("minValue")).isEqualTo(1L); assertThat(partitions.get("partition0").getLong("maxValue")).isEqualTo(7L); assertThat(partitions.get("partition1").getLong("minValue")).isEqualTo(8L); assertThat(partitions.get("partition1").getLong("maxValue")).isEqualTo(14L); assertThat(partitions.get("partition2").getLong("minValue")).isEqualTo(15L); assertThat(partitions.get("partition2").getLong("maxValue")).isEqualTo(21L); assertThat(partitions.get("partition3").getLong("minValue")).isEqualTo(22L); assertThat(partitions.get("partition3").getLong("maxValue")).isEqualTo(28L); assertThat(partitions.get("partition4").getLong("minValue")).isEqualTo(29L); assertThat(partitions.get("partition4").getLong("maxValue")).isEqualTo(31L); }
@Test public void given100Employees_whenPartitioningFor1NodeWith5PartitionsPerNode_thenDistributeCorrectly() throws Exception { List<Long> employeeIds = new ArrayList<>(); for (long i = 1; i <= 100; i++) { employeeIds.add(i); } when(taxCalculationRepository.getUnprocessedEmployeeIds(anyLong(), anyLong(), anyLong())).thenReturn(employeeIds); Map<String, ExecutionContext> partitions = employeeJobPartitioner.partition(1); assertThat(partitions).hasSize(5); assertThat(partitions.get("partition0").getLong("minValue")).isEqualTo(1L); assertThat(partitions.get("partition0").getLong("maxValue")).isEqualTo(21L); assertThat(partitions.get("partition1").getLong("minValue")).isEqualTo(22L); assertThat(partitions.get("partition1").getLong("maxValue")).isEqualTo(42L); assertThat(partitions.get("partition2").getLong("minValue")).isEqualTo(43L); assertThat(partitions.get("partition2").getLong("maxValue")).isEqualTo(63L); assertThat(partitions.get("partition3").getLong("minValue")).isEqualTo(64L); assertThat(partitions.get("partition3").getLong("maxValue")).isEqualTo(84L); assertThat(partitions.get("partition4").getLong("minValue")).isEqualTo(85L); assertThat(partitions.get("partition4").getLong("maxValue")).isEqualTo(100L); }