fix panic of randw when 0 weight

smallnest · Nov 1, 2020 · 55bf8fd · 55bf8fd
1 parent 10873b4
commit 55bf8fd
Show file tree

Hide file tree

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

diff --git a/random_weighted_test.go b/random_weighted_test.go
@@ -54,3 +54,10 @@ func TestRandW_Next(t *testing.T) {
 func checkResults(v, min, max int) bool {
 	return v >= min && v <= max
 }
+
+func TestRandWZero(t *testing.T)  {
+	w := NewRandW()
+	w.Add("a", 0)
+	// test panic or not
+	_ = w.Next()
+}