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

Code before:
            // Verify that the reassignment is running.  The very low throttle should keep it
            // from completing before this runs.
            waitForVerifyAssignment(admin, assignment, true,
                    new VerifyAssignmentResult(partStates, true, Map.of(), false));
            // Cancel the reassignment.
            assertEquals(Map.entry(Set.of(foo0, baz1), Set.of()), runCancelAssignment(assignment, true, useBootstrapServer));
            // Broker throttles are still active because we passed --preserve-throttles
            waitForInterBrokerThrottle(admin, List.of(0, 1, 2, 3), interBrokerThrottle);
            // Cancelling the reassignment again should reveal nothing to cancel.
            assertEquals(Map.entry(Set.of(), Set.of()), runCancelAssignment(assignment, false, useBootstrapServer));
            // This time, the broker throttles were removed.
            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
            // Verify that there are no ongoing reassignments.
            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
        }
        // Verify that the partition is removed from cancelled replicas
        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
    }

    /**
     * Remove a set of throttled partitions and reset the overall replication quota.
     */
    private void removeReplicationThrottleForPartitions(TopicPartition part) {
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            removePartitionReplicaThrottles(admin, new HashSet<>(Set.of(part)));
            assertDoesNotThrow(() -> throttleAllBrokersReplication(admin, List.of(0, 1, 2, 3), Integer.MAX_VALUE));
        }
    }

    private void removePartitionReplicaThrottles(Admin adminClient, Set<TopicPartition> partitions) {
        Map<ConfigResource, Collection<AlterConfigOp>> throttles = partitions.stream()
                .map(tp -> {
                    ConfigResource resource = new ConfigResource(TOPIC, tp.topic());
                    return new AbstractMap.SimpleEntry<>(
                            resource,
                            List.of(
                                    new AlterConfigOp(new ConfigEntry(LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG, ""),
                                            AlterConfigOp.OpType.DELETE),
                                    new AlterConfigOp(new ConfigEntry(FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG, ""),
                                            AlterConfigOp.OpType.DELETE)
                            )
                    );
                })
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
        assertDoesNotThrow(() -> adminClient.incrementalAlterConfigs(throttles).all().get());
    }

    /**
     * Set broker replication quotas and enable throttling for a set of partitions. This
     * will override any previous replication quotas, but will leave the throttling status

Review comment (suggestion):
```suggestion
            removePartitionReplicaThrottles(admin, Set.of(part));
```

Code after:
            // Verify that the reassignment is running.  The very low throttle should keep it
            // from completing before this runs.
            waitForVerifyAssignment(admin, assignment, true,
                    new VerifyAssignmentResult(partStates, true, Map.of(), false));
            // Cancel the reassignment.
            assertEquals(Map.entry(Set.of(foo0, baz1), Set.of()), runCancelAssignment(assignment, true, useBootstrapServer));
            // Broker throttles are still active because we passed --preserve-throttles
            waitForInterBrokerThrottle(admin, List.of(0, 1, 2, 3), interBrokerThrottle);
            // Cancelling the reassignment again should reveal nothing to cancel.
            assertEquals(Map.entry(Set.of(), Set.of()), runCancelAssignment(assignment, false, useBootstrapServer));
            // This time, the broker throttles were removed.
            waitForBrokerLevelThrottles(admin, unthrottledBrokerConfigs);
            // Verify that there are no ongoing reassignments.
            assertFalse(runVerifyAssignment(admin, assignment, false).partsOngoing);
        }
        // Verify that the partition is removed from cancelled replicas
        verifyReplicaDeleted(new TopicPartitionReplica(foo0.topic(), foo0.partition(), 3));
        verifyReplicaDeleted(new TopicPartitionReplica(baz1.topic(), baz1.partition(), 3));
    }

    /**
     * Remove a set of throttled partitions and reset the overall replication quota.
     */
    private void removeReplicationThrottleForPartitions(TopicPartition part) {
        try (Admin admin = Admin.create(Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, clusterInstance.bootstrapServers()))) {
            removePartitionReplicaThrottles(admin, Set.of(part));
            assertDoesNotThrow(() -> throttleAllBrokersReplication(admin, List.of(0, 1, 2, 3), Integer.MAX_VALUE));
        }
    }

    private void removePartitionReplicaThrottles(Admin adminClient, Set<TopicPartition> partitions) {
        Map<ConfigResource, Collection<AlterConfigOp>> throttles = partitions.stream()
                .map(tp -> {
                    ConfigResource resource = new ConfigResource(TOPIC, tp.topic());
                    return new AbstractMap.SimpleEntry<>(
                            resource,
                            List.of(
                                    new AlterConfigOp(new ConfigEntry(LEADER_REPLICATION_THROTTLED_REPLICAS_CONFIG, ""),
                                            AlterConfigOp.OpType.DELETE),
                                    new AlterConfigOp(new ConfigEntry(FOLLOWER_REPLICATION_THROTTLED_REPLICAS_CONFIG, ""),
                                            AlterConfigOp.OpType.DELETE)
                            )
                    );
                })
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
        assertDoesNotThrow(() -> adminClient.incrementalAlterConfigs(throttles).all().get());
    }

    /**
     * Set broker replication quotas and enable throttling for a set of partitions. This
     * will override any previous replication quotas, but will leave the throttling status