github.com/linapex/ethereum-go-chinese@v0.0.0-20190316121929-f8b7a73c3fa1/les/randselect_test.go (about)

     1  
     2  //<developer>
     3  //    <name>linapex 曹一峰</name>
     4  //    <email>linapex@163.com</email>
     5  //    <wx>superexc</wx>
     6  //    <qqgroup>128148617</qqgroup>
     7  //    <url>https://jsq.ink</url>
     8  //    <role>pku engineer</role>
     9  //    <date>2019-03-16 19:16:39</date>
    10  //</624450095974191104>
    11  
    12  
    13  package les
    14  
    15  import (
    16  	"math/rand"
    17  	"testing"
    18  )
    19  
    20  type testWrsItem struct {
    21  	idx  int
    22  	widx *int
    23  }
    24  
    25  func (t *testWrsItem) Weight() int64 {
    26  	w := *t.widx
    27  	if w == -1 || w == t.idx {
    28  		return int64(t.idx + 1)
    29  	}
    30  	return 0
    31  }
    32  
    33  func TestWeightedRandomSelect(t *testing.T) {
    34  	testFn := func(cnt int) {
    35  		s := newWeightedRandomSelect()
    36  		w := -1
    37  		list := make([]testWrsItem, cnt)
    38  		for i := range list {
    39  			list[i] = testWrsItem{idx: i, widx: &w}
    40  			s.update(&list[i])
    41  		}
    42  		w = rand.Intn(cnt)
    43  		c := s.choose()
    44  		if c == nil {
    45  			t.Errorf("expected item, got nil")
    46  		} else {
    47  			if c.(*testWrsItem).idx != w {
    48  				t.Errorf("expected another item")
    49  			}
    50  		}
    51  		w = -2
    52  		if s.choose() != nil {
    53  			t.Errorf("expected nil, got item")
    54  		}
    55  	}
    56  	testFn(1)
    57  	testFn(10)
    58  	testFn(100)
    59  	testFn(1000)
    60  	testFn(10000)
    61  	testFn(100000)
    62  	testFn(1000000)
    63  }
    64