@Override
  public Map<String, Set<Integer>> assignBrokerSetsForUnresolvedBrokers(final ClusterModel clusterModel,
                                                                        final Map<String, Set<Integer>> existingBrokerSetMapping) {
    Set<Broker> allMappedBrokers = existingBrokerSetMapping.values()
                                                           .stream()
                                                           .flatMap(brokerIds -> brokerIds.stream())
                                                           .map(brokerId -> clusterModel.broker(brokerId))
                                                           .filter(Objects::nonNull)
                                                           .collect(Collectors.toSet());

    Set<Broker> extraBrokersInClusterModel = new HashSet<>(clusterModel.brokers());
    extraBrokersInClusterModel.removeAll(allMappedBrokers);

    int numberOfBrokerSets = existingBrokerSetMapping.size();
    List<String> brokerSetIds = new ArrayList<>(existingBrokerSetMapping.keySet());
    Collections.sort(brokerSetIds);

    extraBrokersInClusterModel.stream().forEach(broker -> {
      String brokerSet = brokerSetIds.get(broker.id() % numberOfBrokerSets);
      Set<Integer> brokerIdsForBrokerSet = existingBrokerSetMapping.getOrDefault(brokerSet, new HashSet<>());
      brokerIdsForBrokerSet.add(broker.id());
      existingBrokerSetMapping.put(brokerSet, brokerIdsForBrokerSet);
    });

    return existingBrokerSetMapping;
  }
@Test
  public void testBrokerSetAssignment() {
    // The cluster model has 6 brokers - 0,1,2,3,4,5
    Map<Integer, String> brokers = Map.of(0, "", 1, "", 2, "", 3, "", 4, "", 5, "");
    Map<String, Set<Integer>> mappedBrokers = new HashMap<>();
    mappedBrokers.put("bs1", new HashSet<>(Arrays.asList(0)));
    mappedBrokers.put("bs2", new HashSet<>(Arrays.asList(1)));

    Map<String, Set<Integer>> mappedBrokersAfterAssignment =
        MODULO_BASED_BROKER_SET_ASSIGNMENT_POLICY.assignBrokerSetsForUnresolvedBrokers(brokers, mappedBrokers);

    assertNotNull(mappedBrokersAfterAssignment);

    assertEquals(mappedBrokers.get("bs1"), Set.of(0, 2, 4));
    assertEquals(mappedBrokers.get("bs2"), Set.of(1, 3, 5));
  }