Add unit tests for TaskQueue Upsert and Remove methods

This commit is contained in:
rcourtman 2025-11-29 21:38:33 +00:00
parent 02d8d8a9f9
commit 0dd65c5db2

View file

@ -231,3 +231,572 @@ func TestTaskQueue_Snapshot(t *testing.T) {
})
}
}
func TestTaskQueue_Upsert(t *testing.T) {
t.Run("insert into empty queue", func(t *testing.T) {
queue := NewTaskQueue()
task := ScheduledTask{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: time.Now().Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
}
queue.Upsert(task)
if queue.Size() != 1 {
t.Errorf("Size() = %d, want 1", queue.Size())
}
verifyHeapInvariant(t, queue)
})
t.Run("upsert existing entry with different NextRun", func(t *testing.T) {
queue := NewTaskQueue()
task1 := ScheduledTask{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: time.Now().Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
}
queue.Upsert(task1)
// Update same task with different NextRun
task2 := ScheduledTask{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: time.Now().Add(20 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
}
queue.Upsert(task2)
if queue.Size() != 1 {
t.Errorf("Size() = %d, want 1 (not 2)", queue.Size())
}
// Verify the task was updated
queue.mu.Lock()
key := schedulerKey(InstanceTypePVE, "pve-1")
entry := queue.entries[key]
if !entry.task.NextRun.Equal(task2.NextRun) {
t.Errorf("NextRun not updated: got %v, want %v", entry.task.NextRun, task2.NextRun)
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
t.Run("insert multiple entries - verify heap ordering", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
tasks := []ScheduledTask{
{
InstanceName: "pve-3",
InstanceType: InstanceTypePVE,
NextRun: now.Add(30 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-2",
InstanceType: InstanceTypePVE,
NextRun: now.Add(20 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
if queue.Size() != 3 {
t.Errorf("Size() = %d, want 3", queue.Size())
}
// Verify heap root is earliest NextRun
queue.mu.Lock()
root := queue.heap[0]
if root.task.InstanceName != "pve-1" {
t.Errorf("heap root = %s, want pve-1 (earliest NextRun)", root.task.InstanceName)
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
t.Run("upsert changes heap position - earlier NextRun", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
tasks := []ScheduledTask{
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-2",
InstanceType: InstanceTypePVE,
NextRun: now.Add(20 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-3",
InstanceType: InstanceTypePVE,
NextRun: now.Add(30 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
// Update pve-3 to have earliest NextRun
updatedTask := ScheduledTask{
InstanceName: "pve-3",
InstanceType: InstanceTypePVE,
NextRun: now.Add(5 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
}
queue.Upsert(updatedTask)
if queue.Size() != 3 {
t.Errorf("Size() = %d, want 3", queue.Size())
}
// Verify pve-3 is now at heap root
queue.mu.Lock()
root := queue.heap[0]
if root.task.InstanceName != "pve-3" {
t.Errorf("heap root = %s, want pve-3 (updated to earliest)", root.task.InstanceName)
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
t.Run("upsert changes heap position - later NextRun", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
tasks := []ScheduledTask{
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-2",
InstanceType: InstanceTypePVE,
NextRun: now.Add(20 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-3",
InstanceType: InstanceTypePVE,
NextRun: now.Add(30 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
// Update pve-1 to have latest NextRun
updatedTask := ScheduledTask{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(40 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
}
queue.Upsert(updatedTask)
if queue.Size() != 3 {
t.Errorf("Size() = %d, want 3", queue.Size())
}
// Verify pve-2 is now at heap root
queue.mu.Lock()
root := queue.heap[0]
if root.task.InstanceName != "pve-2" {
t.Errorf("heap root = %s, want pve-2 (pve-1 moved down)", root.task.InstanceName)
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
t.Run("upsert with priority ordering", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
// Same NextRun, different priorities
tasks := []ScheduledTask{
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 0.5,
},
{
InstanceName: "pve-2",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
// Higher priority should be at root
queue.mu.Lock()
root := queue.heap[0]
if root.task.InstanceName != "pve-2" {
t.Errorf("heap root = %s, want pve-2 (higher priority)", root.task.InstanceName)
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
}
func TestTaskQueue_Remove(t *testing.T) {
t.Run("remove from empty queue", func(t *testing.T) {
queue := NewTaskQueue()
// Should not panic
queue.Remove(InstanceTypePVE, "pve-1")
if queue.Size() != 0 {
t.Errorf("Size() = %d, want 0", queue.Size())
}
})
t.Run("remove non-existent key", func(t *testing.T) {
queue := NewTaskQueue()
task := ScheduledTask{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: time.Now().Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
}
queue.Upsert(task)
// Remove different instance
queue.Remove(InstanceTypePVE, "pve-2")
if queue.Size() != 1 {
t.Errorf("Size() = %d, want 1 (pve-1 should still exist)", queue.Size())
}
verifyHeapInvariant(t, queue)
})
t.Run("remove only entry", func(t *testing.T) {
queue := NewTaskQueue()
task := ScheduledTask{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: time.Now().Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
}
queue.Upsert(task)
queue.Remove(InstanceTypePVE, "pve-1")
if queue.Size() != 0 {
t.Errorf("Size() = %d, want 0", queue.Size())
}
queue.mu.Lock()
if len(queue.entries) != 0 {
t.Errorf("entries map has %d items, want 0", len(queue.entries))
}
if len(queue.heap) != 0 {
t.Errorf("heap has %d items, want 0", len(queue.heap))
}
queue.mu.Unlock()
})
t.Run("remove entry from middle of queue", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
tasks := []ScheduledTask{
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-2",
InstanceType: InstanceTypePVE,
NextRun: now.Add(20 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-3",
InstanceType: InstanceTypePVE,
NextRun: now.Add(30 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
// Remove middle entry
queue.Remove(InstanceTypePVE, "pve-2")
if queue.Size() != 2 {
t.Errorf("Size() = %d, want 2", queue.Size())
}
// Verify pve-2 is not in entries map
queue.mu.Lock()
key := schedulerKey(InstanceTypePVE, "pve-2")
if _, exists := queue.entries[key]; exists {
t.Errorf("pve-2 still exists in entries map")
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
t.Run("remove heap root", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
tasks := []ScheduledTask{
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-2",
InstanceType: InstanceTypePVE,
NextRun: now.Add(20 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-3",
InstanceType: InstanceTypePVE,
NextRun: now.Add(30 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
// Remove root (pve-1 has earliest NextRun)
queue.Remove(InstanceTypePVE, "pve-1")
if queue.Size() != 2 {
t.Errorf("Size() = %d, want 2", queue.Size())
}
// Verify pve-2 is now the root
queue.mu.Lock()
root := queue.heap[0]
if root.task.InstanceName != "pve-2" {
t.Errorf("heap root = %s, want pve-2 (next earliest)", root.task.InstanceName)
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
t.Run("remove all entries sequentially", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
tasks := []ScheduledTask{
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-2",
InstanceType: InstanceTypePVE,
NextRun: now.Add(20 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pve-3",
InstanceType: InstanceTypePVE,
NextRun: now.Add(30 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
// Remove each entry and verify invariant after each removal
queue.Remove(InstanceTypePVE, "pve-2")
if queue.Size() != 2 {
t.Errorf("Size after removing pve-2 = %d, want 2", queue.Size())
}
verifyHeapInvariant(t, queue)
queue.Remove(InstanceTypePVE, "pve-1")
if queue.Size() != 1 {
t.Errorf("Size after removing pve-1 = %d, want 1", queue.Size())
}
verifyHeapInvariant(t, queue)
queue.Remove(InstanceTypePVE, "pve-3")
if queue.Size() != 0 {
t.Errorf("Size after removing pve-3 = %d, want 0", queue.Size())
}
verifyHeapInvariant(t, queue)
})
t.Run("remove different instance types", func(t *testing.T) {
queue := NewTaskQueue()
now := time.Now()
tasks := []ScheduledTask{
{
InstanceName: "pve-1",
InstanceType: InstanceTypePVE,
NextRun: now.Add(10 * time.Second),
Interval: 30 * time.Second,
Priority: 1.0,
},
{
InstanceName: "pbs-1",
InstanceType: InstanceTypePBS,
NextRun: now.Add(20 * time.Second),
Interval: 60 * time.Second,
Priority: 1.0,
},
}
for _, task := range tasks {
queue.Upsert(task)
}
// Remove PBS instance
queue.Remove(InstanceTypePBS, "pbs-1")
if queue.Size() != 1 {
t.Errorf("Size() = %d, want 1", queue.Size())
}
// Verify PVE instance still exists
queue.mu.Lock()
key := schedulerKey(InstanceTypePVE, "pve-1")
if _, exists := queue.entries[key]; !exists {
t.Errorf("pve-1 should still exist in entries map")
}
queue.mu.Unlock()
verifyHeapInvariant(t, queue)
})
}
// verifyHeapInvariant checks that the heap maintains its invariants:
// 1. len(entries) matches heap size
// 2. Each entry's index matches its actual position in heap
// 3. Heap property: parent is less than or equal to children
func verifyHeapInvariant(t *testing.T, queue *TaskQueue) {
t.Helper()
queue.mu.Lock()
defer queue.mu.Unlock()
// Check entries count matches heap size
if len(queue.entries) != len(queue.heap) {
t.Errorf("entries count %d != heap size %d", len(queue.entries), len(queue.heap))
}
// Check each entry's index is correct
for _, entry := range queue.heap {
if entry.index < 0 || entry.index >= len(queue.heap) {
t.Errorf("entry %s has invalid index %d (heap size: %d)", entry.key(), entry.index, len(queue.heap))
continue
}
if queue.heap[entry.index] != entry {
t.Errorf("entry %s has index %d but is not at that position in heap", entry.key(), entry.index)
}
}
// Check all entries in map are also in heap
for key, entry := range queue.entries {
found := false
for _, heapEntry := range queue.heap {
if heapEntry == entry {
found = true
break
}
}
if !found {
t.Errorf("entry %s in map but not in heap", key)
}
}
// Check heap property: parent <= children
for i := 0; i < len(queue.heap); i++ {
leftChild := 2*i + 1
rightChild := 2*i + 2
if leftChild < len(queue.heap) {
if queue.heap.Less(leftChild, i) {
t.Errorf("heap violation: child at %d is less than parent at %d", leftChild, i)
}
}
if rightChild < len(queue.heap) {
if queue.heap.Less(rightChild, i) {
t.Errorf("heap violation: child at %d is less than parent at %d", rightChild, i)
}
}
}
}