@Test
public void testGetWeightedRandom() {

  int i;
  float[] weights =
      new float[] {0, 0.1f, 0.2f, 0.2f, -0.1f, 0.1f, 0.2f, 0.1f, 0.1f};
  float[] expectedWeights =
      new float[] {0, 0.1f, 0.2f, 0.2f, 0, 0.1f, 0.2f, 0.1f, 0.1f};
  int[] result = new int[weights.length];
  ArrayList<Float> weightsList = new ArrayList<>();
  for (float weight : weights) {
    weightsList.add(weight);
  }
  int n = 10000000;
  for (i = 0; i < n; i++) {
    int sample = FederationPolicyUtils.getWeightedRandom(weightsList);
    result[sample]++;
  }
  for (i = 0; i < weights.length; i++) {
    double actualWeight = (float) result[i] / n;
    System.out.println(i + " " + actualWeight);
    Assert.assertTrue(
        "Index " + i + " Actual weight: " + actualWeight
            + " expected weight: " + expectedWeights[i],
        Math.abs(actualWeight - expectedWeights[i]) < 0.01);
  }
}