github.com/lyft/flytestdlib@v0.3.12-0.20210213045714-8cdd111ecda1/random/weighted_random_list_test.go (about)

     1  package random
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"strings"
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  type testData struct {
    13  	key string
    14  	val int
    15  }
    16  
    17  func (t testData) Compare(to Comparable) bool {
    18  	if strings.Contains(t.key, "sort") {
    19  		return t.key < to.(testData).key
    20  	}
    21  	return t.val < to.(testData).val
    22  }
    23  
    24  func TestDeterministicWeightedRandomStr(t *testing.T) {
    25  	item1 := testData{
    26  		key: "sort_key1",
    27  		val: 1,
    28  	}
    29  	item2 := testData{
    30  		key: "sort_key2",
    31  		val: 2,
    32  	}
    33  	entries := []Entry{
    34  		{
    35  			Item:   item1,
    36  			Weight: 0.4,
    37  		},
    38  		{
    39  			Item:   item2,
    40  			Weight: 0.6,
    41  		},
    42  	}
    43  	randWeight, err := NewWeightedRandom(context.Background(), entries)
    44  	assert.Nil(t, err)
    45  	retItem, err := randWeight.GetWithSeed(rand.NewSource(20))
    46  	assert.Nil(t, err)
    47  	assert.Equal(t, item1, retItem)
    48  
    49  	assert.Nil(t, err)
    50  	for i := 1; i <= 10; i++ {
    51  		retItem, err := randWeight.GetWithSeed(rand.NewSource(10))
    52  		assert.Nil(t, err)
    53  		assert.Equal(t, item2, retItem)
    54  	}
    55  }
    56  
    57  func TestDeterministicWeightedRandomInt(t *testing.T) {
    58  	item1 := testData{
    59  		key: "key1",
    60  		val: 4,
    61  	}
    62  	item2 := testData{
    63  		key: "key2",
    64  		val: 3,
    65  	}
    66  	entries := []Entry{
    67  		{
    68  			Item:   item1,
    69  			Weight: 0.4,
    70  		},
    71  		{
    72  			Item:   item2,
    73  			Weight: 0.6,
    74  		},
    75  	}
    76  	randWeight, err := NewWeightedRandom(context.Background(), entries)
    77  	assert.Nil(t, err)
    78  	rand.NewSource(10)
    79  	retItem, err := randWeight.GetWithSeed(rand.NewSource(20))
    80  	assert.Nil(t, err)
    81  	assert.Equal(t, item2, retItem)
    82  
    83  	for i := 1; i <= 10; i++ {
    84  		retItem, err := randWeight.GetWithSeed(rand.NewSource(1))
    85  		assert.Nil(t, err)
    86  		assert.Equal(t, item1, retItem)
    87  	}
    88  }
    89  
    90  func TestDeterministicWeightedFewZeroWeight(t *testing.T) {
    91  	item1 := testData{
    92  		key: "key1",
    93  		val: 4,
    94  	}
    95  	item2 := testData{
    96  		key: "key2",
    97  		val: 3,
    98  	}
    99  	entries := []Entry{
   100  		{
   101  			Item:   item1,
   102  			Weight: 0.4,
   103  		},
   104  		{
   105  			Item: item2,
   106  		},
   107  	}
   108  	randWeight, err := NewWeightedRandom(context.Background(), entries)
   109  	assert.Nil(t, err)
   110  	retItem, err := randWeight.GetWithSeed(rand.NewSource(20))
   111  	assert.Nil(t, err)
   112  	assert.Equal(t, item1, retItem)
   113  
   114  	for i := 1; i <= 10; i++ {
   115  		retItem, err := randWeight.GetWithSeed(rand.NewSource(10))
   116  		assert.Nil(t, err)
   117  		assert.Equal(t, item1, retItem)
   118  	}
   119  }
   120  
   121  func TestDeterministicWeightedAllZeroWeights(t *testing.T) {
   122  	item1 := testData{
   123  		key: "sort_key1",
   124  		val: 4,
   125  	}
   126  	item2 := testData{
   127  		key: "sort_key2",
   128  		val: 3,
   129  	}
   130  	entries := []Entry{
   131  		{
   132  			Item: item1,
   133  		},
   134  		{
   135  			Item: item2,
   136  		},
   137  	}
   138  	randWeight, err := NewWeightedRandom(context.Background(), entries)
   139  	assert.Nil(t, err)
   140  	retItem, err := randWeight.GetWithSeed(rand.NewSource(10))
   141  	assert.Nil(t, err)
   142  	assert.Equal(t, item2, retItem)
   143  
   144  	for i := 1; i <= 10; i++ {
   145  		retItem, err := randWeight.GetWithSeed(rand.NewSource(20))
   146  		assert.Nil(t, err)
   147  		assert.Equal(t, item1, retItem)
   148  	}
   149  }
   150  
   151  func TestDeterministicWeightList(t *testing.T) {
   152  	item1 := testData{
   153  		key: "key1",
   154  		val: 4,
   155  	}
   156  	item2 := testData{
   157  		key: "key2",
   158  		val: 3,
   159  	}
   160  	entries := []Entry{
   161  		{
   162  			Item:   item1,
   163  			Weight: 0.3,
   164  		},
   165  		{
   166  			Item: item2,
   167  		},
   168  	}
   169  	randWeight, err := NewWeightedRandom(context.Background(), entries)
   170  	assert.Nil(t, err)
   171  	assert.EqualValues(t, []Comparable{item1}, randWeight.List())
   172  }
   173  
   174  func TestDeterministicWeightListZeroWeights(t *testing.T) {
   175  	item1 := testData{
   176  		key: "key1",
   177  		val: 4,
   178  	}
   179  	item2 := testData{
   180  		key: "key2",
   181  		val: 3,
   182  	}
   183  	entries := []Entry{
   184  		{
   185  			Item: item1,
   186  		},
   187  		{
   188  			Item: item2,
   189  		},
   190  	}
   191  	randWeight, err := NewWeightedRandom(context.Background(), entries)
   192  	assert.Nil(t, err)
   193  	assert.EqualValues(t, []Comparable{item2, item1}, randWeight.List())
   194  }
   195  
   196  func TestDeterministicWeightLen(t *testing.T) {
   197  	item1 := testData{
   198  		key: "key1",
   199  	}
   200  	item2 := testData{
   201  		key: "key2",
   202  	}
   203  	entries := []Entry{
   204  		{
   205  			Item: item1,
   206  		},
   207  		{
   208  			Item: item2,
   209  		},
   210  	}
   211  	randWeight, err := NewWeightedRandom(context.Background(), entries)
   212  	assert.Nil(t, err)
   213  	assert.EqualValues(t, 2, randWeight.Len())
   214  }
   215  
   216  func TestDeterministicWeightInvalidWeights(t *testing.T) {
   217  	item1 := testData{
   218  		key: "key1",
   219  		val: 4,
   220  	}
   221  	item2 := testData{
   222  		key: "key2",
   223  		val: 3,
   224  	}
   225  	entries := []Entry{
   226  		{
   227  			Item:   item1,
   228  			Weight: -3.0,
   229  		},
   230  		{
   231  			Item: item2,
   232  		},
   233  	}
   234  	_, err := NewWeightedRandom(context.Background(), entries)
   235  	assert.NotNil(t, err)
   236  	assert.EqualError(t, err, "invalid weight -3.000000, index 0")
   237  }
   238  
   239  func TestDeterministicWeightInvalidList(t *testing.T) {
   240  	item2 := testData{
   241  		key: "key2",
   242  		val: 3,
   243  	}
   244  	entries := []Entry{
   245  		{},
   246  		{
   247  			Item: item2,
   248  		},
   249  	}
   250  	_, err := NewWeightedRandom(context.Background(), entries)
   251  	assert.NotNil(t, err)
   252  	assert.EqualError(t, err, "invalid entry: nil, index 0")
   253  }