Repository: apache/kafka
File: tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java (Java)

Code before:
                            topicPartition, new PartitionReassignmentState(List.of(0, 1, 2), List.of(0, 1, 2), true)
                    ), false, Map.of(
                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
                            new CompletedMoveState(reassignment.targetDir)
                    ), false));

            BrokerDirs info1 = new BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
        }
    }

    @ClusterTest(types = {Type.KRAFT})
    public void testAlterLogDirReassignmentThrottle() throws Exception {
        createTopics();
        TopicPartition topicPartition = new TopicPartition("foo", 0);
        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);

        int targetBrokerId = 0;
        List<Integer> replicas = List.of(0, 1, 2);
        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);

        // Start the replica move with a low throttle so it does not complete
        long initialLogDirThrottle = 1L;
        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            waitForLogDirThrottle(admin, new HashSet<>(List.of(0)), initialLogDirThrottle);

            // Now increase the throttle and verify that the log dir movement completes
            long updatedLogDirThrottle = 3000000L;
            runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
            waitForLogDirThrottle(admin, Set.of(0), updatedLogDirThrottle);

            waitForVerifyAssignment(admin, reassignment.json, true,
                    new VerifyAssignmentResult(Map.of(
                            topicPartition, new PartitionReassignmentState(List.of(0, 1, 2), List.of(0, 1, 2), true)
                    ), false, Map.of(
                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
                            new CompletedMoveState(reassignment.targetDir)
                    ), false));
        }
    }

    @ClusterTest
    public void testDisallowReplicationFactorChange() {
        createTopics();
        String assignment = "{\"version\":1,\"partitions\":" +
                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1],\"log_dirs\":[\"any\",\"any\"]}," +
                "{\"topic\":\"foo\",\"partition\":1,\"replicas\":[0,1,2,3],\"log_dirs\":[\"any\",\"any\",\"any\",\"any\"]}," +
                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3],\"log_dirs\":[\"any\"]}" +
                "]}";
        try (Admin admin = clusterInstance.admin()) {

Review comment (suggestion):
```suggestion
            waitForLogDirThrottle(admin, Set.of(0), initialLogDirThrottle);
```

Code after:
                            topicPartition, new PartitionReassignmentState(List.of(0, 1, 2), List.of(0, 1, 2), true)
                    ), false, Map.of(
                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), 0),
                            new CompletedMoveState(reassignment.targetDir)
                    ), false));

            BrokerDirs info1 = new BrokerDirs(admin.describeLogDirs(IntStream.range(0, 4).boxed().collect(Collectors.toList())), 0);
            assertEquals(reassignment.targetDir, info1.curLogDirs.getOrDefault(topicPartition, ""));
        }
    }

    @ClusterTest(types = {Type.KRAFT})
    public void testAlterLogDirReassignmentThrottle() throws Exception {
        createTopics();
        TopicPartition topicPartition = new TopicPartition("foo", 0);
        produceMessages(topicPartition.topic(), topicPartition.partition(), 700);

        int targetBrokerId = 0;
        List<Integer> replicas = List.of(0, 1, 2);
        LogDirReassignment reassignment = buildLogDirReassignment(topicPartition, targetBrokerId, replicas);

        // Start the replica move with a low throttle so it does not complete
        long initialLogDirThrottle = 1L;
        runExecuteAssignment(false, reassignment.json, -1L, initialLogDirThrottle);
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            waitForLogDirThrottle(admin, Set.of(0), initialLogDirThrottle);

            // Now increase the throttle and verify that the log dir movement completes
            long updatedLogDirThrottle = 3000000L;
            runExecuteAssignment(true, reassignment.json, -1L, updatedLogDirThrottle);
            waitForLogDirThrottle(admin, Set.of(0), updatedLogDirThrottle);

            waitForVerifyAssignment(admin, reassignment.json, true,
                    new VerifyAssignmentResult(Map.of(
                            topicPartition, new PartitionReassignmentState(List.of(0, 1, 2), List.of(0, 1, 2), true)
                    ), false, Map.of(
                            new TopicPartitionReplica(topicPartition.topic(), topicPartition.partition(), targetBrokerId),
                            new CompletedMoveState(reassignment.targetDir)
                    ), false));
        }
    }

    @ClusterTest
    public void testDisallowReplicationFactorChange() {
        createTopics();
        String assignment = "{\"version\":1,\"partitions\":" +
                "[{\"topic\":\"foo\",\"partition\":0,\"replicas\":[0,1],\"log_dirs\":[\"any\",\"any\"]}," +
                "{\"topic\":\"foo\",\"partition\":1,\"replicas\":[0,1,2,3],\"log_dirs\":[\"any\",\"any\",\"any\",\"any\"]}," +
                "{\"topic\":\"bar\",\"partition\":0,\"replicas\":[3],\"log_dirs\":[\"any\"]}" +
                "]}";
        try (Admin admin = clusterInstance.admin()) {