github.com/nsqio/nsq@v1.3.0/internal/pqueue/pqueue_test.go (about)

     1  package pqueue
     2  
     3  import (
     4  	"container/heap"
     5  	"math/rand"
     6  	"path/filepath"
     7  	"reflect"
     8  	"runtime"
     9  	"sort"
    10  	"testing"
    11  )
    12  
    13  func equal(t *testing.T, act, exp interface{}) {
    14  	if !reflect.DeepEqual(exp, act) {
    15  		_, file, line, _ := runtime.Caller(1)
    16  		t.Logf("\033[31m%s:%d:\n\n\texp: %#v\n\n\tgot: %#v\033[39m\n\n",
    17  			filepath.Base(file), line, exp, act)
    18  		t.FailNow()
    19  	}
    20  }
    21  
    22  func TestPriorityQueue(t *testing.T) {
    23  	c := 100
    24  	pq := New(c)
    25  
    26  	for i := 0; i < c+1; i++ {
    27  		heap.Push(&pq, &Item{Value: i, Priority: int64(i)})
    28  	}
    29  	equal(t, pq.Len(), c+1)
    30  	equal(t, cap(pq), c*2)
    31  
    32  	for i := 0; i < c+1; i++ {
    33  		item := heap.Pop(&pq)
    34  		equal(t, item.(*Item).Value.(int), i)
    35  	}
    36  	equal(t, cap(pq), c/4)
    37  }
    38  
    39  func TestUnsortedInsert(t *testing.T) {
    40  	c := 100
    41  	pq := New(c)
    42  	ints := make([]int, 0, c)
    43  
    44  	for i := 0; i < c; i++ {
    45  		v := rand.Int()
    46  		ints = append(ints, v)
    47  		heap.Push(&pq, &Item{Value: i, Priority: int64(v)})
    48  	}
    49  	equal(t, pq.Len(), c)
    50  	equal(t, cap(pq), c)
    51  
    52  	sort.Ints(ints)
    53  
    54  	for i := 0; i < c; i++ {
    55  		item, _ := pq.PeekAndShift(int64(ints[len(ints)-1]))
    56  		equal(t, item.Priority, int64(ints[i]))
    57  	}
    58  }
    59  
    60  func TestRemove(t *testing.T) {
    61  	c := 100
    62  	pq := New(c)
    63  
    64  	for i := 0; i < c; i++ {
    65  		v := rand.Int()
    66  		heap.Push(&pq, &Item{Value: "test", Priority: int64(v)})
    67  	}
    68  
    69  	for i := 0; i < 10; i++ {
    70  		heap.Remove(&pq, rand.Intn((c-1)-i))
    71  	}
    72  
    73  	lastPriority := heap.Pop(&pq).(*Item).Priority
    74  	for i := 0; i < (c - 10 - 1); i++ {
    75  		item := heap.Pop(&pq)
    76  		equal(t, lastPriority < item.(*Item).Priority, true)
    77  		lastPriority = item.(*Item).Priority
    78  	}
    79  }