agentic-coding-workflow/internal/orchestrator/scheduler_test.go

package orchestrator

import (
	"context"
	"sync"
	"testing"

	"github.com/google/uuid"
	"github.com/stretchr/testify/require"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	"github.com/yan1h/agent-coding-workflow/internal/project"
)

// ===== fakeReadySource: orchestrator.ReadyTaskSource =====

type fakeReadySource struct {
	mu        sync.Mutex
	projects  []uuid.UUID
	ready     map[uuid.UUID][]*project.Issue // projectID → ready issues
	active    map[uuid.UUID]int              // projectID → active session count
	listCalls int
}

func (f *fakeReadySource) ListSchedulableProjects(_ context.Context) ([]uuid.UUID, error) {
	return f.projects, nil
}

func (f *fakeReadySource) ListReadyLeafSubtasks(_ context.Context, projectID uuid.UUID, limit int) ([]*project.Issue, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.listCalls++
	all := f.ready[projectID]
	// 模拟 DB：排除已被启动（active）的 issue 由调用方维护；此处仅按 limit 截断。
	if limit > 0 && len(all) > limit {
		return all[:limit], nil
	}
	return all, nil
}

func (f *fakeReadySource) CountActiveSessionsForProject(_ context.Context, projectID uuid.UUID) (int, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.active[projectID], nil
}

// ===== fakeSessionStarter: orchestrator.SessionStarter =====

type fakeSessionStarter struct {
	mu        sync.Mutex
	started   []uuid.UUID // issue ids
	quotaFrom int         // 第 quotaFrom 次起返回配额错误（0 表示不触发）
	calls     int
	failAll   bool
}

func (f *fakeSessionStarter) StartForIssue(_ context.Context, iss *project.Issue) (uuid.UUID, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.calls++
	if f.failAll {
		return uuid.Nil, errs.New(errs.CodeInternal, "boom")
	}
	if f.quotaFrom > 0 && f.calls >= f.quotaFrom {
		return uuid.Nil, errs.New(errs.CodeAcpSessionQuotaExceeded, "quota")
	}
	f.started = append(f.started, iss.ID)
	return uuid.New(), nil
}

// ===== fakeSchedNotify: orchestrator.SchedulerNotifier =====

type fakeSchedNotify struct {
	dispatched int
	quota      int
}

func (f *fakeSchedNotify) NotifyDispatched(_ context.Context, _ *project.Issue, _ uuid.UUID) {
	f.dispatched++
}
func (f *fakeSchedNotify) NotifyQuotaBlocked(_ context.Context, _ *project.Issue) { f.quota++ }

func mkReadyIssues(projectID uuid.UUID, n int) []*project.Issue {
	out := make([]*project.Issue, 0, n)
	for i := 0; i < n; i++ {
		out = append(out, &project.Issue{ID: uuid.New(), ProjectID: projectID, Number: i + 1})
	}
	return out
}

func TestScheduler_DispatchesUpToFanout(t *testing.T) {
	projID := uuid.New()
	src := &fakeReadySource{
		projects: []uuid.UUID{projID},
		ready:    map[uuid.UUID][]*project.Issue{projID: mkReadyIssues(projID, 5)},
		active:   map[uuid.UUID]int{},
	}
	starter := &fakeSessionStarter{}
	notify := &fakeSchedNotify{}
	sched := NewScheduler(SchedulerDeps{
		Projects: src, Sessions: starter, Notify: notify,
		MaxFanoutPerTick: 3,
	})

	n, err := sched.Tick(context.Background())
	require.NoError(t, err)
	require.Equal(t, 3, n, "应只派发 min(ready=5, fanout=3)")
	require.Len(t, starter.started, 3)
	require.Equal(t, 3, notify.dispatched)
}

func TestScheduler_QuotaIsSoftFailure(t *testing.T) {
	projID := uuid.New()
	src := &fakeReadySource{
		projects: []uuid.UUID{projID},
		ready:    map[uuid.UUID][]*project.Issue{projID: mkReadyIssues(projID, 5)},
		active:   map[uuid.UUID]int{},
	}
	// 第 3 次起返回配额错误：前 2 个成功，第 3 个软失败 → 停止本 project。
	starter := &fakeSessionStarter{quotaFrom: 3}
	notify := &fakeSchedNotify{}
	sched := NewScheduler(SchedulerDeps{
		Projects: src, Sessions: starter, Notify: notify,
		MaxFanoutPerTick: 5,
	})

	n, err := sched.Tick(context.Background())
	require.NoError(t, err, "配额超限是软失败，不应返回 error")
	require.Equal(t, 2, n)
	require.Len(t, starter.started, 2)
	require.Equal(t, 1, notify.quota)
}

func TestScheduler_OtherErrorSkipsIssueNotProject(t *testing.T) {
	projID := uuid.New()
	src := &fakeReadySource{
		projects: []uuid.UUID{projID},
		ready:    map[uuid.UUID][]*project.Issue{projID: mkReadyIssues(projID, 3)},
		active:   map[uuid.UUID]int{},
	}
	starter := &fakeSessionStarter{failAll: true}
	sched := NewScheduler(SchedulerDeps{
		Projects: src, Sessions: starter, MaxFanoutPerTick: 5,
	})
	n, err := sched.Tick(context.Background())
	require.NoError(t, err)
	require.Equal(t, 0, n, "全部 start 失败但非配额错误：跳过 issue，整轮不报错")
}

func TestScheduler_MaxConcurrentPerProjectPreCheck(t *testing.T) {
	projID := uuid.New()
	src := &fakeReadySource{
		projects: []uuid.UUID{projID},
		ready:    map[uuid.UUID][]*project.Issue{projID: mkReadyIssues(projID, 10)},
		active:   map[uuid.UUID]int{projID: 3}, // 已有 3 个活跃
	}
	starter := &fakeSessionStarter{}
	sched := NewScheduler(SchedulerDeps{
		Projects: src, Sessions: starter,
		MaxFanoutPerTick:        10,
		MaxConcurrentPerProject: 4, // budget = 4-3 = 1
	})
	n, err := sched.Tick(context.Background())
	require.NoError(t, err)
	require.Equal(t, 1, n, "软并发预检：budget=1，本 tick 只派 1 个")
}

func TestScheduler_AtConcurrencyCapDispatchesNothing(t *testing.T) {
	projID := uuid.New()
	src := &fakeReadySource{
		projects: []uuid.UUID{projID},
		ready:    map[uuid.UUID][]*project.Issue{projID: mkReadyIssues(projID, 5)},
		active:   map[uuid.UUID]int{projID: 4},
	}
	starter := &fakeSessionStarter{}
	sched := NewScheduler(SchedulerDeps{
		Projects: src, Sessions: starter,
		MaxFanoutPerTick: 10, MaxConcurrentPerProject: 4,
	})
	n, err := sched.Tick(context.Background())
	require.NoError(t, err)
	require.Equal(t, 0, n)
	require.Empty(t, starter.started)
}

func TestScheduler_IdempotentAcrossTicks(t *testing.T) {
	projID := uuid.New()
	issues := mkReadyIssues(projID, 2)
	src := &fakeReadySource{
		projects: []uuid.UUID{projID},
		ready:    map[uuid.UUID][]*project.Issue{projID: issues},
		active:   map[uuid.UUID]int{},
	}
	starter := &fakeSessionStarter{}
	sched := NewScheduler(SchedulerDeps{Projects: src, Sessions: starter, MaxFanoutPerTick: 5})

	n1, err := sched.Tick(context.Background())
	require.NoError(t, err)
	require.Equal(t, 2, n1)

	// 第二 tick：模拟 DB 已排除被启动的 issue（清空 ready），不应重复派发。
	src.mu.Lock()
	src.ready[projID] = nil
	src.mu.Unlock()
	n2, err := sched.Tick(context.Background())
	require.NoError(t, err)
	require.Equal(t, 0, n2)
	require.Len(t, starter.started, 2, "跨 tick 不重复派发")
}