freedakgmail
111 lines
2.7 KiB
Go
111 lines
2.7 KiB
Go
package scheduler
|
|
|
|
import (
|
|
"context"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/edgeai/gateway/internal/config"
|
|
"github.com/edgeai/gateway/internal/observability"
|
|
"github.com/edgeai/gateway/internal/task"
|
|
)
|
|
|
|
func newTestScheduler(maxRunning, maxQueued int) *Scheduler {
|
|
cfg := &config.SchedulerConfig{
|
|
MaxRunningTasks: maxRunning,
|
|
MaxQueuedTasks: maxQueued,
|
|
}
|
|
logger := observability.NewLogger(observability.LevelDebug, os.Stdout, "metadata_only")
|
|
return NewScheduler(cfg, logger)
|
|
}
|
|
|
|
func TestSubmitAndGetNext(t *testing.T) {
|
|
s := newTestScheduler(2, 10)
|
|
defer s.Stop()
|
|
|
|
task1 := task.NewTask("t1", "r1", "app1", "tenant1", "model1", task.PriorityNormal, false)
|
|
task2 := task.NewTask("t2", "r2", "app1", "tenant1", "model1", task.PriorityHigh, false)
|
|
|
|
if err := s.Submit(task1); err != nil {
|
|
t.Fatalf("submit task1: %v", err)
|
|
}
|
|
if err := s.Submit(task2); err != nil {
|
|
t.Fatalf("submit task2: %v", err)
|
|
}
|
|
|
|
ctx := context.Background()
|
|
got1, err := s.GetNext(ctx)
|
|
if err != nil {
|
|
t.Fatalf("get next: %v", err)
|
|
}
|
|
// P1 (High) should come before P2 (Normal)
|
|
if got1.ID != "t2" {
|
|
t.Errorf("expected t2 (higher priority) first, got %s", got1.ID)
|
|
}
|
|
|
|
got2, err := s.GetNext(ctx)
|
|
if err != nil {
|
|
t.Fatalf("get next 2: %v", err)
|
|
}
|
|
if got2.ID != "t1" {
|
|
t.Errorf("expected t1 second, got %s", got2.ID)
|
|
}
|
|
}
|
|
|
|
func TestQueueFull(t *testing.T) {
|
|
s := newTestScheduler(1, 2)
|
|
defer s.Stop()
|
|
|
|
for i := 0; i < 2; i++ {
|
|
tk := task.NewTask("t", "r", "app", "tenant", "model", task.PriorityNormal, false)
|
|
if err := s.Submit(tk); err != nil {
|
|
t.Fatalf("submit %d: %v", i, err)
|
|
}
|
|
}
|
|
|
|
tk := task.NewTask("t3", "r3", "app", "tenant", "model", task.PriorityNormal, false)
|
|
err := s.Submit(tk)
|
|
if err == nil {
|
|
t.Error("expected queue full error")
|
|
}
|
|
}
|
|
|
|
func TestComplete(t *testing.T) {
|
|
s := newTestScheduler(1, 10)
|
|
defer s.Stop()
|
|
|
|
tk := task.NewTask("t1", "r1", "app", "tenant", "model", task.PriorityNormal, false)
|
|
s.Submit(tk)
|
|
|
|
ctx := context.Background()
|
|
got, _ := s.GetNext(ctx)
|
|
if s.RunningCount() != 1 {
|
|
t.Errorf("expected 1 running, got %d", s.RunningCount())
|
|
}
|
|
|
|
s.Complete(got.ID)
|
|
if s.RunningCount() != 0 {
|
|
t.Errorf("expected 0 running after complete, got %d", s.RunningCount())
|
|
}
|
|
}
|
|
|
|
func TestFIFOOrdering(t *testing.T) {
|
|
s := newTestScheduler(1, 10)
|
|
defer s.Stop()
|
|
|
|
// Same priority, should be FIFO
|
|
t1 := task.NewTask("t1", "r1", "app", "tenant", "model", task.PriorityNormal, false)
|
|
time.Sleep(1 * time.Millisecond)
|
|
t2 := task.NewTask("t2", "r2", "app", "tenant", "model", task.PriorityNormal, false)
|
|
|
|
s.Submit(t1)
|
|
s.Submit(t2)
|
|
|
|
ctx := context.Background()
|
|
got1, _ := s.GetNext(ctx)
|
|
if got1.ID != "t1" {
|
|
t.Errorf("expected t1 first (FIFO), got %s", got1.ID)
|
|
}
|
|
}
|