github.com/sberex/go-sberex@v1.8.2-0.20181113200658-ed96ac38f7d7/les/randselect_test.go (about) 1 // This file is part of the go-sberex library. The go-sberex library is 2 // free software: you can redistribute it and/or modify it under the terms 3 // of the GNU Lesser General Public License as published by the Free 4 // Software Foundation, either version 3 of the License, or (at your option) 5 // any later version. 6 // 7 // The go-sberex library is distributed in the hope that it will be useful, 8 // but WITHOUT ANY WARRANTY; without even the implied warranty of 9 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser 10 // General Public License <http://www.gnu.org/licenses/> for more details. 11 12 package les 13 14 import ( 15 "math/rand" 16 "testing" 17 ) 18 19 type testWrsItem struct { 20 idx int 21 widx *int 22 } 23 24 func (t *testWrsItem) Weight() int64 { 25 w := *t.widx 26 if w == -1 || w == t.idx { 27 return int64(t.idx + 1) 28 } 29 return 0 30 } 31 32 func TestWeightedRandomSelect(t *testing.T) { 33 testFn := func(cnt int) { 34 s := newWeightedRandomSelect() 35 w := -1 36 list := make([]testWrsItem, cnt) 37 for i := range list { 38 list[i] = testWrsItem{idx: i, widx: &w} 39 s.update(&list[i]) 40 } 41 w = rand.Intn(cnt) 42 c := s.choose() 43 if c == nil { 44 t.Errorf("expected item, got nil") 45 } else { 46 if c.(*testWrsItem).idx != w { 47 t.Errorf("expected another item") 48 } 49 } 50 w = -2 51 if s.choose() != nil { 52 t.Errorf("expected nil, got item") 53 } 54 } 55 testFn(1) 56 testFn(10) 57 testFn(100) 58 testFn(1000) 59 testFn(10000) 60 testFn(100000) 61 testFn(1000000) 62 }