github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/priority_queue_test.go (about)

     1  package util
     2  
     3  import (
     4  	"runtime"
     5  	"strconv"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  )
    11  
    12  type simpleItem int64
    13  
    14  func (i simpleItem) Priority() int64 {
    15  	return int64(i)
    16  }
    17  
    18  func (i simpleItem) Key() string {
    19  	return strconv.FormatInt(int64(i), 10)
    20  }
    21  
    22  func TestPriorityQueueBasic(t *testing.T) {
    23  	queue := NewPriorityQueue(nil)
    24  	assert.Equal(t, 0, queue.Length(), "Expected length = 0")
    25  
    26  	queue.Enqueue(simpleItem(1))
    27  	assert.Equal(t, 1, queue.Length(), "Expected length = 1")
    28  
    29  	i, ok := queue.Dequeue().(simpleItem)
    30  	assert.True(t, ok, "Expected cast to succeed")
    31  	assert.Equal(t, simpleItem(1), i, "Expected to dequeue simpleItem(1)")
    32  
    33  	queue.Close()
    34  	assert.Nil(t, queue.Dequeue(), "Expect nil dequeue")
    35  }
    36  
    37  func TestPriorityQueuePriorities(t *testing.T) {
    38  	queue := NewPriorityQueue(nil)
    39  	queue.Enqueue(simpleItem(1))
    40  	queue.Enqueue(simpleItem(2))
    41  
    42  	assert.Equal(t, simpleItem(2), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(2)")
    43  	assert.Equal(t, simpleItem(1), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(1)")
    44  
    45  	queue.Close()
    46  	assert.Nil(t, queue.Dequeue(), "Expect nil dequeue")
    47  }
    48  
    49  func TestPriorityQueuePriorities2(t *testing.T) {
    50  	queue := NewPriorityQueue(nil)
    51  	queue.Enqueue(simpleItem(2))
    52  	queue.Enqueue(simpleItem(1))
    53  
    54  	assert.Equal(t, simpleItem(2), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(2)")
    55  	assert.Equal(t, simpleItem(1), queue.Dequeue().(simpleItem), "Expected to dequeue simpleItem(1)")
    56  
    57  	queue.Close()
    58  	assert.Nil(t, queue.Dequeue(), "Expect nil dequeue")
    59  }
    60  
    61  func TestPriorityQueueWait(t *testing.T) {
    62  	queue := NewPriorityQueue(nil)
    63  
    64  	done := make(chan struct{})
    65  	go func() {
    66  		assert.Nil(t, queue.Dequeue(), "Expect nil dequeue")
    67  		close(done)
    68  	}()
    69  
    70  	queue.Close()
    71  	runtime.Gosched()
    72  	select {
    73  	case <-done:
    74  	case <-time.After(100 * time.Millisecond):
    75  		t.Fatal("Close didn't unblock Dequeue.")
    76  	}
    77  }