Merge pull request #9 from linhan-cai/bug/rand_number_fix

bug fix RandW
smallnest · Aug 20, 2020 · 10873b4 · 10873b4
2 parents adf21c9 + 60fe40a
commit 10873b4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/random_weighted.go b/random_weighted.go
@@ -30,7 +30,7 @@ func (rw *RandW) Next() (item interface{}) {
 	if rw.n == 0 {
 		return nil
 	}
-	randomWeight := rw.r.Intn(rw.sumOfWeights)
+	randomWeight := rw.r.Intn(rw.sumOfWeights) + 1
 	for _, item := range rw.items {
 		randomWeight = randomWeight - item.Weight
 		if randomWeight <= 0 {

diff --git a/random_weighted_test.go b/random_weighted_test.go
@@ -4,9 +4,10 @@ import "testing"
 
 func TestRandW_Next(t *testing.T) {
 	w := NewRandW()
-	w.Add("server1", 5)
+	w.Add("server1", 4)
 	w.Add("server2", 2)
 	w.Add("server3", 3)
+	w.Add("server4", 1)
 
 	results := make(map[string]int)
 
@@ -15,7 +16,7 @@ func TestRandW_Next(t *testing.T) {
 		results[s]++
 	}
 
-	if !checkResults(results["server1"], 20, 70) || !checkResults(results["server2"], 0, 40) || !checkResults(results["server3"], 10, 50) {
+	if !checkResults(results["server1"], 20, 70) || !checkResults(results["server2"], 1, 40) || !checkResults(results["server3"], 10, 50) || !checkResults(results["server4"], 1, 30) {
 		t.Error("the algorithm is wrong", results)
 	}
 
@@ -27,7 +28,7 @@ func TestRandW_Next(t *testing.T) {
 		results[s]++
 	}
 
-	if !checkResults(results["server1"], 20, 70) || !checkResults(results["server2"], 0, 40) || !checkResults(results["server3"], 10, 50) {
+	if !checkResults(results["server1"], 20, 70) || !checkResults(results["server2"], 1, 40) || !checkResults(results["server3"], 10, 50) || !checkResults(results["server4"], 1, 30) {
 		t.Error("the algorithm is wrong", results)
 	}