boolean canSwap(Replica r1, Replica r2, ClusterModel clusterModel) {
    String mappedRackIdOfR1 = mappedRackIdOf(r1.broker());
    String mappedRackIdOfR2 = mappedRackIdOf(r2.broker());
    boolean inSameRack = r1.broker() != r2.broker() && mappedRackIdOfR1.equals(mappedRackIdOfR2);
    boolean rackAware =
        clusterModel.partition(r1.topicPartition()).partitionRacks().stream()
                    .map(this::mappedRackIdOf)
                    .noneMatch(mappedRackIdOfR2::equals)
        && clusterModel.partition(r2.topicPartition()).partitionRacks().stream()
                       .map(this::mappedRackIdOf)
                       .noneMatch(mappedRackIdOfR1::equals);

    boolean sameRole = r1.isLeader() == r2.isLeader();
    return (inSameRack || rackAware) && sameRole;
  }
@Test
  public void testCanSwap() {
    KafkaAssignerDiskUsageDistributionGoal goal = new KafkaAssignerDiskUsageDistributionGoal();
    ClusterModel clusterModel = createClusterModel();

    Replica r1 = clusterModel.broker(0).replica(T0P0);
    Replica r2 = clusterModel.broker(1).replica(T2P0);
    assertTrue("Replicas in the same rack should be good to swap", goal.canSwap(r1, r2, clusterModel));
    assertTrue("Replicas in the same rack should be good to swap", goal.canSwap(r2, r1, clusterModel));

    r2 = clusterModel.broker(1).replica(T1P0);
    assertFalse("Should not be able to swap replica with different roles.", goal.canSwap(r1, r2, clusterModel));
    assertFalse("Should not be able to swap replica with different roles.", goal.canSwap(r2, r1, clusterModel));

    r2 = clusterModel.broker(2).replica(T2P1);
    assertFalse("Should not be able to put two replicas in the same broker", goal.canSwap(r1, r2, clusterModel));
    assertFalse("Should not be able to put two replicas in the same broker", goal.canSwap(r2, r1, clusterModel));

    r2 = clusterModel.broker(3).replica(T2P2);
    assertFalse("Should not be able to put two replicas in the same rack", goal.canSwap(r1, r2, clusterModel));
    assertFalse("Should not be able to put two replicas in the same rack", goal.canSwap(r2, r1, clusterModel));

    r1 = clusterModel.broker(3).replica(T0P2);
    r2 = clusterModel.broker(4).replica(T1P2);
    assertTrue("Should be able to swap", goal.canSwap(r1, r2, clusterModel));
    assertTrue("Should be able to swap", goal.canSwap(r2, r1, clusterModel));
  }