go-zero/core/stat/topk_test.go

package stat

import (
	"math/rand"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
)

const (
	numSamples = 10000
	topNum     = 100
)

var samples []Task

func init() {
	for i := 0; i < numSamples; i++ {
		task := Task{
			Duration: time.Duration(rand.Int63()),
		}
		samples = append(samples, task)
	}
}

func TestTopK(t *testing.T) {
	tasks := []Task{
		{false, 1, "a"},
		{false, 4, "a"},
		{false, 2, "a"},
		{false, 5, "a"},
		{false, 9, "a"},
		{false, 10, "a"},
		{false, 12, "a"},
		{false, 3, "a"},
		{false, 6, "a"},
		{false, 11, "a"},
		{false, 8, "a"},
	}

	result := topK(tasks, 3)
	if len(result) != 3 {
		t.Fail()
	}

	set := make(map[time.Duration]struct{})
	for _, each := range result {
		set[each.Duration] = struct{}{}
	}

	for _, v := range []time.Duration{10, 11, 12} {
		_, ok := set[v]
		assert.True(t, ok)
	}
}

func BenchmarkTopkHeap(b *testing.B) {
	for i := 0; i < b.N; i++ {
		topK(samples, topNum)
	}
}