prysm-pulse/container/queue/priority_queue_test.go

package queue

import (
	"container/heap"
	"fmt"
	"testing"
	"time"

	"github.com/prysmaticlabs/prysm/v5/testing/require"
)

// Ensure we satisfy the heap.Interface
var _ heap.Interface = &queue{}

// some tests rely on the ordering of items from this method
func testCases() (tc []*Item) {
	// create a slice of items with priority / times offset by these seconds
	for i, m := range []time.Duration{
		5,
		183600,  // 51 hours
		15,      // 15 seconds
		45,      // 45 seconds
		900,     // 15 minutes
		300,     // 5 minutes
		7200,    // 2 hours
		183600,  // 51 hours
		7201,    // 2 hours, 1 second
		115200,  // 32 hours
		1209600, // 2 weeks
	} {
		n := time.Now()
		ft := n.Add(time.Second * m)
		tc = append(tc, &Item{
			Key:      fmt.Sprintf("item-%d", i),
			Value:    1,
			Priority: ft.Unix(),
		})
	}
	return
}

func TestPriorityQueue_New(t *testing.T) {
	pq := New()

	if len(pq.data) != len(pq.dataMap) || len(pq.data) != 0 {
		t.Fatalf("error in queue/map size, expected data and map to be initialized, got (%d) and (%d)", len(pq.data), len(pq.dataMap))
	}

	if pq.Len() != 0 {
		t.Fatalf("expected new queue to have zero size, got (%d)", pq.Len())
	}
}

func TestPriorityQueue_Push(t *testing.T) {
	pq := New()

	// don't allow nil pushing
	if err := pq.Push(nil); err == nil {
		t.Fatal("Expected error on pushing nil")
	}

	tc := testCases()
	tcl := len(tc)
	for _, i := range tc {
		if err := pq.Push(i); err != nil {
			t.Fatal(err)
		}
	}

	if pq.Len() != tcl {
		t.Fatalf("error adding items, expected (%d) items, got (%d)", tcl, pq.Len())
	}

	testValidateInternalData(t, pq, len(tc), false)

	item, err := pq.Pop()
	if err != nil {
		t.Fatalf("error popping item: %s", err)
	}
	if tc[0].Priority != item.Priority {
		t.Fatalf("expected tc[0] and popped item to match, got (%q) and (%q)", tc[0], item.Priority)
	}
	if tc[0].Key != item.Key {
		t.Fatalf("expected tc[0] and popped item to match, got (%q) and (%q)", tc[0], item.Priority)
	}

	testValidateInternalData(t, pq, len(tc)-1, false)

	// push item with no key
	dErr := pq.Push(tc[1])
	if dErr != ErrDuplicateItem {
		t.Fatal(err)
	}
	// push item with no key
	tc[2].Key = ""
	kErr := pq.Push(tc[2])
	if kErr != nil && kErr.Error() != "error adding item: Item Key is required" {
		t.Fatal(kErr)
	}

	testValidateInternalData(t, pq, len(tc)-1, false)

	// check nil,nil error for not found
	i, err := pq.PopByKey("empty")
	if err != nil && i != nil {
		t.Fatalf("expected nil error for PopByKey of non-existing key, got: %s", err)
	}
}

func TestPriorityQueue_Pop(t *testing.T) {
	pq := New()

	tc := testCases()
	for _, i := range tc {
		if err := pq.Push(i); err != nil {
			t.Fatal(err)
		}
	}

	topItem, err := pq.Pop()
	if err != nil {
		t.Fatalf("error calling pop: %s", err)
	}
	if tc[0].Priority != topItem.Priority {
		t.Fatalf("expected tc[0] and popped item to match, got (%q) and (%q)", tc[0], topItem.Priority)
	}
	if tc[0].Key != topItem.Key {
		t.Fatalf("expected tc[0] and popped item to match, got (%q) and (%q)", tc[0], topItem.Priority)
	}

	var items []*Item
	items = append(items, topItem)
	// pop the remaining items, compare size of input and output
	i, err := pq.Pop()
	require.NoError(t, err)
	for ; i != nil; i, err = pq.Pop() {
		require.NoError(t, err)
		items = append(items, i)
	}
	require.Equal(t, len(tc), len(items))
}

func TestPriorityQueue_PopByKey(t *testing.T) {
	pq := New()

	tc := testCases()
	for _, i := range tc {
		if err := pq.Push(i); err != nil {
			t.Fatal(err)
		}
	}

	// grab the top priority item, to capture it's value for checking later
	item, err := pq.Pop()
	require.NoError(t, err)
	oldPriority := item.Priority
	oldKey := item.Key

	// push the item back on, so it gets removed with PopByKey and we verify
	// the top item has changed later
	require.NoError(t, pq.Push(item))

	popKeys := []int{2, 4, 7, 1, 0}
	for _, i := range popKeys {
		item, err := pq.PopByKey(fmt.Sprintf("item-%d", i))
		if err != nil {
			t.Fatalf("failed to pop item-%d, \n\terr: %s\n\titem: %#v", i, err, item)
		}
	}

	testValidateInternalData(t, pq, len(tc)-len(popKeys), false)

	// grab the top priority item again, to compare with the top item priority
	// from above
	item, err = pq.Pop()
	require.NoError(t, err)
	newPriority := item.Priority
	newKey := item.Key

	if oldPriority == newPriority || oldKey == newKey {
		t.Fatalf("expected old/new key and priority to differ, got (%s/%s) and (%d/%d)", oldKey, newKey, oldPriority, newPriority)
	}

	testValidateInternalData(t, pq, len(tc)-len(popKeys)-1, true)
}

func TestPriorityQueue_RetrieveByKey(t *testing.T) {
	pq := New()

	tc := testCases()
	for _, i := range tc {
		if err := pq.Push(i); err != nil {
			t.Fatal(err)
		}
	}

	// grab the top priority item, to capture it's value for checking later
	item, err := pq.Pop()
	require.NoError(t, err)
	oldPriority := item.Priority
	oldKey := item.Key

	// push the item back on, so it gets retrieved with RetrieveByKey and we verify
	// the top item does not change.
	require.NoError(t, pq.Push(item))

	popKeys := []int{2, 4, 7, 1, 0}
	for _, i := range popKeys {
		item := pq.RetrieveByKey(fmt.Sprintf("item-%d", i))
		if item == nil {
			t.Fatalf("failed to pop item-%d, \n\titem: %#v", i, item)
		}
	}

	testValidateInternalData(t, pq, len(tc), false)

	// grab the top priority item again, to compare with the top item priority
	// from above. They should be the same
	item, err = pq.Pop()
	require.NoError(t, err)
	newPriority := item.Priority
	newKey := item.Key

	if oldPriority != newPriority && oldKey != newKey {
		t.Fatalf("expected old/new key and priority to be the same, got (%s/%s) and (%d/%d)", oldKey, newKey, oldPriority, newPriority)
	}

	testValidateInternalData(t, pq, len(tc)-1, true)

	// ensure the correct item is retrieved when multiple items exist
	require.NoError(t, pq.Push(&Item{
		Key:   "baz",
		Value: "foo",
	}))
	require.NoError(t, pq.Push(&Item{
		Key:   "foo",
		Value: "bar",
	}))
	i := pq.RetrieveByKey("foo")
	require.Equal(t, "bar", i.Value)

}

// testValidateInternalData checks the internal data structure of the PriorityQueue
// and verifies that items are in-sync. Use drain only at the end of a test,
// because it will mutate the input queue
func testValidateInternalData(t *testing.T, pq *PriorityQueue, expectedSize int, drain bool) {
	actualSize := pq.Len()
	if actualSize != expectedSize {
		t.Fatalf("expected new queue size to be (%d), got (%d)", expectedSize, actualSize)
	}

	if len(pq.data) != len(pq.dataMap) || len(pq.data) != expectedSize {
		t.Fatalf("error in queue/map size, expected data and map to be (%d), got (%d) and (%d)", expectedSize, len(pq.data), len(pq.dataMap))
	}

	if drain && pq.Len() > 0 {
		// pop all the items, verify lengths
		i, err := pq.Pop()
		require.NoError(t, err)
		for ; i != nil; i, err = pq.Pop() {
			require.NoError(t, err)
			expectedSize--
			if len(pq.data) != len(pq.dataMap) || len(pq.data) != expectedSize {
				t.Fatalf("error in queue/map size, expected data and map to be (%d), got (%d) and (%d)", expectedSize, len(pq.data), len(pq.dataMap))
			}
		}
	}
}