Examples with DescribeReplicaLogDirsResult org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult used on opensource projects

Example 1 with DescribeReplicaLogDirsResult

use of org.apache.kafka.clients.admin.DescribeReplicaLogDirsResult in project cruise-control by linkedin.

the class ExecutionTaskPlannerTest method testGetIntraBrokerPartitionMovementTasks.

@Test
public void testGetIntraBrokerPartitionMovementTasks() {
    ReplicaPlacementInfo r0d0 = new ReplicaPlacementInfo(0, "d0");
    ReplicaPlacementInfo r0d1 = new ReplicaPlacementInfo(0, "d1");
    ReplicaPlacementInfo r1d0 = new ReplicaPlacementInfo(1, "d0");
    ReplicaPlacementInfo r1d1 = new ReplicaPlacementInfo(1, "d1");
    List<ExecutionProposal> proposals = Collections.singletonList(new ExecutionProposal(new TopicPartition(TOPIC2, 0), 4, r0d0, Arrays.asList(r0d0, r1d1), Arrays.asList(r1d0, r0d1)));
    TopicPartitionReplica tpr0 = new TopicPartitionReplica(TOPIC2, 0, 0);
    TopicPartitionReplica tpr1 = new TopicPartitionReplica(TOPIC2, 0, 1);
    // Mock adminClient
    AdminClient mockAdminClient = EasyMock.mock(AdminClient.class);
    try {
        // Reflectively set constructors from package private to public.
        Constructor<DescribeReplicaLogDirsResult> constructor1 = DescribeReplicaLogDirsResult.class.getDeclaredConstructor(Map.class);
        constructor1.setAccessible(true);
        Constructor<ReplicaLogDirInfo> constructor2 = ReplicaLogDirInfo.class.getDeclaredConstructor(String.class, long.class, String.class, long.class);
        constructor2.setAccessible(true);
        Map<TopicPartitionReplica, KafkaFuture<ReplicaLogDirInfo>> futureByReplica = new HashMap<>();
        futureByReplica.put(tpr0, completedFuture(constructor2.newInstance("d0", 0L, null, -1L)));
        futureByReplica.put(tpr1, completedFuture(constructor2.newInstance("d1", 0L, null, -1L)));
        EasyMock.expect(mockAdminClient.describeReplicaLogDirs(anyObject())).andReturn(constructor1.newInstance(futureByReplica)).anyTimes();
        EasyMock.replay(mockAdminClient);
    } catch (InstantiationException | IllegalAccessException | InvocationTargetException | NoSuchMethodException e) {
    // Let it go.
    }
    Set<PartitionInfo> partitions = new HashSet<>();
    Node[] isrArray = generateExpectedReplicas(proposals.get(0));
    partitions.add(new PartitionInfo(proposals.get(0).topicPartition().topic(), proposals.get(0).topicPartition().partition(), isrArray[0], isrArray, isrArray));
    Cluster expectedCluster = new Cluster(null, _expectedNodes, partitions, Collections.emptySet(), Collections.emptySet());
    StrategyOptions strategyOptions = new StrategyOptions.Builder(expectedCluster).build();
    ExecutionTaskPlanner planner = new ExecutionTaskPlanner(mockAdminClient, new KafkaCruiseControlConfig(KafkaCruiseControlUnitTestUtils.getKafkaCruiseControlProperties()));
    planner.addExecutionProposals(proposals, strategyOptions, null);
    assertEquals(1, planner.remainingLeadershipMovements().size());
    assertEquals(2, planner.remainingIntraBrokerReplicaMovements().size());
    planner.clear();
    assertEquals(0, planner.remainingLeadershipMovements().size());
    assertEquals(0, planner.remainingIntraBrokerReplicaMovements().size());
    EasyMock.verify(mockAdminClient);
}