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) }