iop/packages/go/agenttask/scheduler_test.go

83 lines
2.5 KiB
Go

package agenttask
import (
"context"
"errors"
"testing"
"time"
)
func TestSchedulerProviderCapacityAndParallelRelease(t *testing.T) {
units := []WorkUnit{
testUnit("one", WriteSetDisjoint),
testUnit("two", WriteSetOverlap),
testUnit("three", WriteSetUnknown),
}
snapshot := testSnapshot("project", "workspace", units...)
harness := newHarness(t, map[ProjectID]ProjectWorkflowSnapshot{"project": snapshot}, 2)
for _, unit := range units {
harness.invoker.delays[unit.ID] = 25 * time.Millisecond
}
harness.start("project", "workspace", nil)
if err := harness.manager.Reconcile(context.Background()); err != nil {
t.Fatalf("Reconcile: %v", err)
}
if maximum := harness.invoker.maxConcurrency(); maximum != 2 {
t.Fatalf("max provider concurrency = %d, want 2", maximum)
}
if harness.invoker.callCount() != 3 {
t.Fatalf("provider calls = %d, want 3", harness.invoker.callCount())
}
}
func TestSchedulerCancelReleasesTicket(t *testing.T) {
scheduler := NewScheduler()
target := ExecutionTarget{
ProviderID: "provider", ProfileID: "profile", Capacity: 1,
}
first, err := scheduler.Acquire(context.Background(), DispatchCandidate{
ProjectID: "p1", WorkspaceID: "w", WorkUnitID: "one", AttemptID: "one#1",
Target: target,
})
if err != nil {
t.Fatalf("first Acquire: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
cancel()
_, err = scheduler.Acquire(ctx, DispatchCandidate{
ProjectID: "p2", WorkspaceID: "w", WorkUnitID: "two", AttemptID: "two#1",
Target: target,
})
if !errors.Is(err, context.Canceled) {
t.Fatalf("cancelled Acquire error = %v", err)
}
first.Release()
second, err := scheduler.Acquire(context.Background(), DispatchCandidate{
ProjectID: "p2", WorkspaceID: "w", WorkUnitID: "two", AttemptID: "two#1",
Target: target,
})
if err != nil {
t.Fatalf("Acquire after release: %v", err)
}
second.Release()
if scheduler.Active(target.PoolKey()) != 0 {
t.Fatalf("scheduler leaked capacity")
}
}
func TestIsolatedDispatchUsesDistinctTaskRoots(t *testing.T) {
snapshot := testSnapshot(
"project", "workspace",
testUnit("overlap-a", WriteSetOverlap),
testUnit("overlap-b", WriteSetOverlap),
)
harness := newHarness(t, map[ProjectID]ProjectWorkflowSnapshot{"project": snapshot}, 2)
harness.start("project", "workspace", nil)
if err := harness.manager.Reconcile(context.Background()); err != nil {
t.Fatalf("Reconcile: %v", err)
}
roots := harness.invoker.roots()
if len(roots) != 2 || roots[0] == roots[1] {
t.Fatalf("isolated task roots = %v", roots)
}
}