@VisibleForTesting
  protected ArrayList<Integer> computeIntegerAssignment(int totalNum,
      ArrayList<Float> weightsList) throws YarnException {
    int i, residue;
    ArrayList<Integer> ret = new ArrayList<>();
    float totalWeight = 0, totalNumFloat = totalNum;

    if (weightsList.size() == 0) {
      return ret;
    }
    for (i = 0; i < weightsList.size(); i++) {
      ret.add(0);
      if (weightsList.get(i) > 0) {
        totalWeight += weightsList.get(i);
      }
    }
    if (totalWeight == 0) {
      StringBuilder sb = new StringBuilder();
      for (Float weight : weightsList) {
        sb.append(weight + ", ");
      }
      throw new FederationPolicyException(
          "No positive value found in weight array " + sb.toString());
    }

    // First pass, do flooring for all bins
    residue = totalNum;
    for (i = 0; i < weightsList.size(); i++) {
      if (weightsList.get(i) > 0) {
        int base = (int) (totalNumFloat * weightsList.get(i) / totalWeight);
        ret.set(i, ret.get(i) + base);
        residue -= base;
      }
    }

    // By now residue < weights.length, assign one a time
    for (i = 0; i < residue; i++) {
      int index = FederationPolicyUtils.getWeightedRandom(weightsList);
      ret.set(index, ret.get(index) + 1);
    }
    return ret;
  }
@Test
  public void testIntegerAssignment() throws YarnException {
    float[] weights =
        new float[] {0, 0.1f, 0.2f, 0.2f, -0.1f, 0.1f, 0.2f, 0.1f, 0.1f};
    int[] expectedMin = new int[] {0, 1, 3, 3, 0, 1, 3, 1, 1};
    ArrayList<Float> weightsList = new ArrayList<>();
    for (float weight : weights) {
      weightsList.add(weight);
    }

    LocalityMulticastAMRMProxyPolicy policy =
        (LocalityMulticastAMRMProxyPolicy) getPolicy();
    for (int i = 0; i < 500000; i++) {
      ArrayList<Integer> allocations =
          policy.computeIntegerAssignment(19, weightsList);
      int sum = 0;
      for (int j = 0; j < weights.length; j++) {
        sum += allocations.get(j);
        if (allocations.get(j) < expectedMin[j]) {
          Assert.fail(allocations.get(j) + " at index " + j
              + " should be at least " + expectedMin[j] + ". Allocation array: "
              + printList(allocations));
        }
      }
      Assert.assertEquals(
          "Expect sum to be 19 in array: " + printList(allocations), 19, sum);
    }
  }