diff --git a/core/external_metadata.go b/core/external_metadata.go index d4efa6b0e..c8efca137 100644 --- a/core/external_metadata.go +++ b/core/external_metadata.go @@ -7,8 +7,6 @@ import ( "sync" "time" - "github.com/navidrome/navidrome/utils" - "github.com/Masterminds/squirrel" "github.com/microcosm-cc/bluemonday" "github.com/navidrome/navidrome/conf" @@ -16,6 +14,7 @@ import ( "github.com/navidrome/navidrome/core/agents" "github.com/navidrome/navidrome/log" "github.com/navidrome/navidrome/model" + "github.com/navidrome/navidrome/utils" ) const ( diff --git a/utils/weighted_random_chooser.go b/utils/weighted_random_chooser.go index 387922240..8692f148a 100644 --- a/utils/weighted_random_chooser.go +++ b/utils/weighted_random_chooser.go @@ -29,6 +29,9 @@ func (w *weightedChooser) Put(value interface{}, weight int) { // GetAndRemove choose a random entry based on their weights, and removes it from the list func (w *weightedChooser) GetAndRemove() (interface{}, error) { + if w.totalWeight == 0 { + return nil, errors.New("cannot choose from zero weight") + } i, err := w.weightedChoice() if err != nil { return nil, err diff --git a/utils/weighted_random_chooser_test.go b/utils/weighted_random_chooser_test.go new file mode 100644 index 000000000..d7303d02c --- /dev/null +++ b/utils/weighted_random_chooser_test.go @@ -0,0 +1,50 @@ +package utils + +import ( + . "github.com/onsi/ginkgo" + . "github.com/onsi/gomega" +) + +var _ = Describe("WeightedRandomChooser", func() { + var w *weightedChooser + BeforeEach(func() { + w = NewWeightedRandomChooser() + for i := 0; i < 10; i++ { + w.Put(i, i) + } + }) + + It("removes a random item", func() { + Expect(w.Size()).To(Equal(10)) + _, err := w.GetAndRemove() + Expect(err).ToNot(HaveOccurred()) + Expect(w.Size()).To(Equal(9)) + }) + + It("returns the sole item", func() { + w = NewWeightedRandomChooser() + w.Put("a", 1) + Expect(w.GetAndRemove()).To(Equal("a")) + }) + + It("fails when trying to choose from empty set", func() { + w = NewWeightedRandomChooser() + w.Put("a", 1) + w.Put("b", 1) + Expect(w.GetAndRemove()).To(BeElementOf("a", "b")) + Expect(w.GetAndRemove()).To(BeElementOf("a", "b")) + _, err := w.GetAndRemove() + Expect(err).To(HaveOccurred()) + }) + + It("chooses based on weights", func() { + counts := [10]int{} + for i := 0; i < 200000; i++ { + c, _ := w.weightedChoice() + counts[c] = counts[c] + 1 + } + for i := 0; i < 9; i++ { + Expect(counts[i]).To(BeNumerically("<", counts[i+1])) + } + }) +})