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  }