github.com/lingyao2333/mo-zero@v1.4.1/core/stat/topk_test.go (about)

     1  package stat
     2  
     3  import (
     4  	"math/rand"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  const (
    12  	numSamples = 10000
    13  	topNum     = 100
    14  )
    15  
    16  var samples []Task
    17  
    18  func init() {
    19  	for i := 0; i < numSamples; i++ {
    20  		task := Task{
    21  			Duration: time.Duration(rand.Int63()),
    22  		}
    23  		samples = append(samples, task)
    24  	}
    25  }
    26  
    27  func TestTopK(t *testing.T) {
    28  	tasks := []Task{
    29  		{false, 1, "a"},
    30  		{false, 4, "a"},
    31  		{false, 2, "a"},
    32  		{false, 5, "a"},
    33  		{false, 9, "a"},
    34  		{false, 10, "a"},
    35  		{false, 12, "a"},
    36  		{false, 3, "a"},
    37  		{false, 6, "a"},
    38  		{false, 11, "a"},
    39  		{false, 8, "a"},
    40  	}
    41  
    42  	result := topK(tasks, 3)
    43  	if len(result) != 3 {
    44  		t.Fail()
    45  	}
    46  
    47  	set := make(map[time.Duration]struct{})
    48  	for _, each := range result {
    49  		set[each.Duration] = struct{}{}
    50  	}
    51  
    52  	for _, v := range []time.Duration{10, 11, 12} {
    53  		_, ok := set[v]
    54  		assert.True(t, ok)
    55  	}
    56  }
    57  
    58  func BenchmarkTopkHeap(b *testing.B) {
    59  	for i := 0; i < b.N; i++ {
    60  		topK(samples, topNum)
    61  	}
    62  }