fetch_ml/internal/scheduler/priority_queue_test.go

package scheduler_test

import (
	"fmt"
	"testing"
	"time"

	"github.com/jfraeys/fetch_ml/internal/scheduler"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestPriorityQueue_BasicOperations(t *testing.T) {
	q := scheduler.NewPriorityQueue(0.1)

	task1 := &scheduler.Task{
		ID:          "task-1",
		Priority:    10,
		SubmittedAt: time.Now(),
		Spec:        scheduler.JobSpec{ID: "task-1"},
	}

	task2 := &scheduler.Task{
		ID:          "task-2",
		Priority:    5,
		SubmittedAt: time.Now(),
		Spec:        scheduler.JobSpec{ID: "task-2"},
	}

	// Add tasks
	q.Add(task1)
	q.Add(task2)

	require.Equal(t, 2, q.Len())

	// Should return highest priority first (task1 with priority 10)
	first := q.Take()
	require.NotNil(t, first)
	assert.Equal(t, "task-1", first.ID)

	// Second task
	second := q.Take()
	require.NotNil(t, second)
	assert.Equal(t, "task-2", second.ID)

	// Queue empty
	third := q.Take()
	assert.Nil(t, third)
}

func TestPriorityQueue_EffectivePriority_WithAging(t *testing.T) {
	now := time.Now()

	// Task with lower priority but older submission
	oldTask := &scheduler.Task{
		ID:          "old-task",
		Priority:    5,
		SubmittedAt: now.Add(-10 * time.Minute), // 10 min old
	}

	// Task with higher priority but recent submission
	newTask := &scheduler.Task{
		ID:          "new-task",
		Priority:    10,
		SubmittedAt: now, // Just submitted
	}

	// Calculate effective priorities
	oldEffective := oldTask.EffectivePriority(0.1, now)
	newEffective := newTask.EffectivePriority(0.1, now)

	// Old task should have higher effective priority due to aging
	// 5 + (10 min * 0.1) = 6.0
	// 10 + (0 min * 0.1) = 10.0
	assert.Less(t, oldEffective, newEffective)
	assert.InDelta(t, 6.0, oldEffective, 0.1)
	assert.InDelta(t, 10.0, newEffective, 0.1)
}

func TestPriorityQueue_FIFOOnTie(t *testing.T) {
	now := time.Now()
	q := scheduler.NewPriorityQueue(0.1)

	// Two tasks with same priority, submitted at different times
	task1 := &scheduler.Task{
		ID:          "task-1",
		Priority:    10,
		SubmittedAt: now.Add(-5 * time.Minute),
		Spec:        scheduler.JobSpec{ID: "task-1"},
	}

	task2 := &scheduler.Task{
		ID:          "task-2",
		Priority:    10,
		SubmittedAt: now.Add(-1 * time.Minute),
		Spec:        scheduler.JobSpec{ID: "task-2"},
	}

	// Add in reverse order
	q.Add(task2)
	q.Add(task1)

	// Should return older task first (FIFO on tie)
	first := q.Take()
	require.NotNil(t, first)
	assert.Equal(t, "task-1", first.ID)

	second := q.Take()
	require.NotNil(t, second)
	assert.Equal(t, "task-2", second.ID)
}

func TestPriorityQueue_Remove(t *testing.T) {
	q := scheduler.NewPriorityQueue(0.1)

	task1 := &scheduler.Task{ID: "task-1", Priority: 10, Spec: scheduler.JobSpec{ID: "task-1"}}
	task2 := &scheduler.Task{ID: "task-2", Priority: 5, Spec: scheduler.JobSpec{ID: "task-2"}}
	task3 := &scheduler.Task{ID: "task-3", Priority: 1, Spec: scheduler.JobSpec{ID: "task-3"}}

	q.Add(task1)
	q.Add(task2)
	q.Add(task3)

	// Remove middle task
	removed := q.Remove("task-2")
	assert.True(t, removed)
	assert.Equal(t, 2, q.Len())

	// Try to remove non-existent
	removed = q.Remove("non-existent")
	assert.False(t, removed)

	// Verify remaining order
	first := q.Take()
	assert.Equal(t, "task-1", first.ID)
	second := q.Take()
	assert.Equal(t, "task-3", second.ID)
}

func TestPriorityQueue_Get(t *testing.T) {
	q := scheduler.NewPriorityQueue(0.1)

	task1 := &scheduler.Task{ID: "task-1", Priority: 10, Spec: scheduler.JobSpec{ID: "task-1"}}
	q.Add(task1)

	// Get existing task
	found := q.Get("task-1")
	assert.NotNil(t, found)
	assert.Equal(t, "task-1", found.ID)

	// Get non-existent
	notFound := q.Get("non-existent")
	assert.Nil(t, notFound)
}

func TestPriorityQueue_Items(t *testing.T) {
	q := scheduler.NewPriorityQueue(0.1)

	tasks := []*scheduler.Task{
		{ID: "task-1", Priority: 10, Spec: scheduler.JobSpec{ID: "task-1"}},
		{ID: "task-2", Priority: 5, Spec: scheduler.JobSpec{ID: "task-2"}},
		{ID: "task-3", Priority: 1, Spec: scheduler.JobSpec{ID: "task-3"}},
	}

	for _, task := range tasks {
		q.Add(task)
	}

	items := q.Items()
	require.Len(t, items, 3)

	// Items should be in priority order (highest first)
	assert.Equal(t, "task-1", items[0].ID)
	assert.Equal(t, "task-2", items[1].ID)
	assert.Equal(t, "task-3", items[2].ID)
}

func TestPriorityQueue_ConcurrentAccess(t *testing.T) {
	q := scheduler.NewPriorityQueue(0.1)
	done := make(chan bool, 3)

	// Concurrent adds
	go func() {
		for i := 0; i < 100; i++ {
			q.Add(&scheduler.Task{ID: fmt.Sprintf("task-%d", i), Priority: i})
		}
		done <- true
	}()

	// Concurrent takes
	go func() {
		for i := 0; i < 50; i++ {
			q.Take()
		}
		done <- true
	}()

	// Concurrent peeks
	go func() {
		for i := 0; i < 100; i++ {
			q.Peek()
		}
		done <- true
	}()

	// Wait for all goroutines
	for i := 0; i < 3; i++ {
		<-done
	}

	// Queue should be in consistent state
	assert.GreaterOrEqual(t, q.Len(), 0)
	assert.LessOrEqual(t, q.Len(), 100)
}