s_cherry_e/sr/internal/extend/rand/weight_test.go

package rand

import (
	"testing"
)

func TestSliceWeight(t *testing.T) {
	list := []int32{1, 2, 3}

	weight := NewSliceWeight[int32](list, func(v int32) int32 {
		return v
	})

	for i := 0; i < 10; i++ {
		t.Log(weight.SelectOne())
	}
}

func TestMapWeight(t *testing.T) {
	type Item struct {
		Key   int32
		Value int32
	}

	maps := map[int32]*Item{
		1: {
			Key:   1,
			Value: 2,
		},
		2: {
			Key:   2,
			Value: 3,
		},
		3: {
			Key:   3,
			Value: 4,
		},
	}

	weight := NewMapWeight[int32](maps, func(v *Item) int32 {
		return v.Value
	})

	for i := 0; i < 10; i++ {
		t.Log(weight.SelectOne())
	}
}