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

package orchestrator

import (
	"context"
	"encoding/json"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"

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

// ===== fakeRepo: 内存版 orchestrator.Repository =====

type fakeRepo struct {
	mu    sync.Mutex
	runs  map[uuid.UUID]*Run
	steps map[uuid.UUID]*Step
}

func newFakeRepo() *fakeRepo {
	return &fakeRepo{runs: map[uuid.UUID]*Run{}, steps: map[uuid.UUID]*Step{}}
}

func cloneRun(r *Run) *Run    { c := *r; return &c }
func cloneStep(s *Step) *Step { c := *s; return &c }

func (r *fakeRepo) InsertRun(_ context.Context, run *Run) (*Run, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if run.ID == uuid.Nil {
		run.ID = uuid.New()
	}
	run.CreatedAt = time.Now()
	run.UpdatedAt = run.CreatedAt
	r.runs[run.ID] = cloneRun(run)
	return cloneRun(run), nil
}

func (r *fakeRepo) GetRun(_ context.Context, id uuid.UUID) (*Run, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	run, ok := r.runs[id]
	if !ok {
		return nil, errs.New(errs.CodeNotFound, "run not found")
	}
	return cloneRun(run), nil
}

func (r *fakeRepo) GetActiveRunByRequirement(_ context.Context, reqID uuid.UUID) (*Run, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, run := range r.runs {
		if run.RequirementID == reqID && run.Status.IsActive() || (run.RequirementID == reqID && run.Status == RunPaused) {
			return cloneRun(run), nil
		}
	}
	return nil, errs.New(errs.CodeNotFound, "no active run")
}

func (r *fakeRepo) ListRuns(_ context.Context, f RunFilter) ([]*Run, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := []*Run{}
	for _, run := range r.runs {
		if f.RequirementID != nil && run.RequirementID != *f.RequirementID {
			continue
		}
		if f.Status != "" && run.Status != f.Status {
			continue
		}
		out = append(out, cloneRun(run))
	}
	return out, nil
}

func (r *fakeRepo) UpdateRunStatus(_ context.Context, id uuid.UUID, status RunStatus, lastErr *string) (*Run, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	run, ok := r.runs[id]
	if !ok {
		return nil, errs.New(errs.CodeNotFound, "run not found")
	}
	run.Status = status
	run.LastError = lastErr
	run.UpdatedAt = time.Now()
	return cloneRun(run), nil
}

func (r *fakeRepo) UpdateRunPhase(_ context.Context, id uuid.UUID, phase project.Phase) (*Run, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	run, ok := r.runs[id]
	if !ok {
		return nil, errs.New(errs.CodeNotFound, "run not found")
	}
	run.CurrentPhase = phase
	run.UpdatedAt = time.Now()
	return cloneRun(run), nil
}

func (r *fakeRepo) InsertStep(_ context.Context, s *Step) (*Step, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if s.ID == uuid.Nil {
		s.ID = uuid.New()
	}
	s.CreatedAt = time.Now()
	s.UpdatedAt = s.CreatedAt
	r.steps[s.ID] = cloneStep(s)
	return cloneStep(s), nil
}

func (r *fakeRepo) GetStep(_ context.Context, id uuid.UUID) (*Step, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	s, ok := r.steps[id]
	if !ok {
		return nil, errs.New(errs.CodeNotFound, "step not found")
	}
	return cloneStep(s), nil
}

func (r *fakeRepo) GetStepBySession(_ context.Context, sessionID uuid.UUID) (*Step, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, s := range r.steps {
		if s.ACPSessionID != nil && *s.ACPSessionID == sessionID {
			return cloneStep(s), nil
		}
	}
	return nil, errs.New(errs.CodeNotFound, "step not found")
}

func (r *fakeRepo) ListStepsByRun(_ context.Context, runID uuid.UUID) ([]*Step, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := []*Step{}
	for _, s := range r.steps {
		if s.RunID == runID {
			out = append(out, cloneStep(s))
		}
	}
	return out, nil
}

func (r *fakeRepo) UpdateStepStatus(_ context.Context, s *Step) (*Step, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	cur, ok := r.steps[s.ID]
	if !ok {
		return nil, errs.New(errs.CodeNotFound, "step not found")
	}
	cur.Status = s.Status
	cur.ACPSessionID = s.ACPSessionID
	cur.StopReason = s.StopReason
	cur.GateResult = s.GateResult
	cur.LastError = s.LastError
	cur.JobID = s.JobID
	cur.UpdatedAt = time.Now()
	return cloneStep(cur), nil
}

func (r *fakeRepo) IncrementStepAttempt(_ context.Context, id uuid.UUID) (*Step, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	s, ok := r.steps[id]
	if !ok {
		return nil, errs.New(errs.CodeNotFound, "step not found")
	}
	s.Attempt++
	s.UpdatedAt = time.Now()
	return cloneStep(s), nil
}

func (r *fakeRepo) CountActiveChildSteps(_ context.Context, parentStepID uuid.UUID) (int64, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	var n int64
	for _, s := range r.steps {
		if s.ParentStepID != nil && *s.ParentStepID == parentStepID {
			switch s.Status {
			case StepPending, StepSpawning, StepRunning, StepAwaitingGate:
				n++
			}
		}
	}
	return n, nil
}

func (r *fakeRepo) ResetStuckStepsOnRestart(_ context.Context) ([]*Step, error) {
	r.mu.Lock()
	defer r.mu.Unlock()
	out := []*Step{}
	for _, s := range r.steps {
		run, ok := r.runs[s.RunID]
		if !ok {
			continue
		}
		active := run.Status.IsActive() || run.Status == RunPaused
		switch s.Status {
		case StepSpawning, StepRunning, StepAwaitingGate:
			if active {
				s.Status = StepPending
				out = append(out, cloneStep(s))
			}
		}
	}
	return out, nil
}

// fakeRepo 不支持真正的事务：InTx 直接调用 fn（同一实例），WithTx 返回自身。
func (r *fakeRepo) InTx(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) error {
	return fn(ctx, nil)
}
func (r *fakeRepo) WithTx(_ pgx.Tx) Repository { return r }

// ===== fakePM: 内存版 PMAccess =====

type fakePM struct {
	mu          sync.Mutex
	proj        *project.Project
	req         *project.Requirement
	artifacts   []*project.Artifact
	phaseWrites []project.Phase
	phaseErr    error
}

func (p *fakePM) GetProjectByID(_ context.Context, _ uuid.UUID) (*project.Project, error) {
	return p.proj, nil
}
func (p *fakePM) GetProjectBySlug(_ context.Context, _ string) (*project.Project, error) {
	return p.proj, nil
}
func (p *fakePM) GetRequirementByID(_ context.Context, _ uuid.UUID) (*project.Requirement, error) {
	return p.req, nil
}
func (p *fakePM) GetRequirementByNumber(_ context.Context, _ uuid.UUID, _ int) (*project.Requirement, error) {
	return p.req, nil
}
func (p *fakePM) ListLatestArtifacts(_ context.Context, _ uuid.UUID) ([]*project.Artifact, error) {
	return p.artifacts, nil
}
func (p *fakePM) ChangeRequirementPhase(_ context.Context, _ uuid.UUID, _ bool, _ string, _ int, to project.Phase) error {
	p.mu.Lock()
	defer p.mu.Unlock()
	if p.phaseErr != nil {
		return p.phaseErr
	}
	p.phaseWrites = append(p.phaseWrites, to)
	if p.req != nil {
		p.req.Phase = to
	}
	return nil
}

// ===== fakeTurnRunner: acp.TurnRunner =====

type fakeTurnRunner struct {
	stopReason string
	err        error
	resumeErr  error
	created    *acp.Session
	resumed    *acp.Session
	calls      int
}

func (f *fakeTurnRunner) SendPromptAndWait(_ context.Context, sessionID uuid.UUID, _ string) (string, error) {
	f.calls++
	return f.stopReason, f.err
}
func (f *fakeTurnRunner) ResumeSession(_ context.Context, _ acp.Caller, sessionID uuid.UUID) (*acp.Session, error) {
	if f.resumeErr != nil {
		return nil, f.resumeErr
	}
	return &acp.Session{ID: sessionID, Status: acp.SessionStarting}, nil
}

// ===== fakeACPSessions: orchestrator.ACPSessions =====

type fakeACPSessions struct {
	mu        sync.Mutex
	created   []acp.CreateSessionInput
	nextID    uuid.UUID
	createErr error
}

func (f *fakeACPSessions) Create(_ context.Context, _ acp.Caller, in acp.CreateSessionInput) (*acp.Session, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	if f.createErr != nil {
		return nil, f.createErr
	}
	f.created = append(f.created, in)
	id := f.nextID
	if id == uuid.Nil {
		id = uuid.New()
	}
	return &acp.Session{ID: id, Status: acp.SessionStarting, OrchestratorStepID: in.OrchestratorStepID}, nil
}

// ===== fakeLookup: orchestrator.ACPSessionLookup =====

type fakeLookup struct {
	crashed map[uuid.UUID]*acp.Session // stepID → crashed session
	live    map[uuid.UUID]*acp.Session // stepID → live session
}

func newFakeLookup() *fakeLookup {
	return &fakeLookup{crashed: map[uuid.UUID]*acp.Session{}, live: map[uuid.UUID]*acp.Session{}}
}
func (f *fakeLookup) GetSessionByStepID(_ context.Context, stepID uuid.UUID) (*acp.Session, error) {
	if s, ok := f.live[stepID]; ok {
		return s, nil
	}
	return nil, errs.New(errs.CodeNotFound, "no session")
}
func (f *fakeLookup) GetCrashedSessionForResume(_ context.Context, stepID uuid.UUID) (*acp.Session, error) {
	if s, ok := f.crashed[stepID]; ok {
		return s, nil
	}
	return nil, errs.New(errs.CodeNotFound, "no crashed session")
}

// ===== fakeJobs: orchestrator.JobEnqueuer =====

type fakeJobs struct {
	mu       sync.Mutex
	enqueued []stepPayload
}

func (j *fakeJobs) Enqueue(_ context.Context, _ jobs.JobType, payload json.RawMessage, _ time.Time, _ int32) (*jobs.Job, error) {
	j.mu.Lock()
	defer j.mu.Unlock()
	var p stepPayload
	_ = json.Unmarshal(payload, &p)
	j.enqueued = append(j.enqueued, p)
	return &jobs.Job{ID: uuid.New()}, nil
}
func (j *fakeJobs) count() int {
	j.mu.Lock()
	defer j.mu.Unlock()
	return len(j.enqueued)
}

// ===== fakeGate: orchestrator.PhaseGate =====

type fakeGate struct {
	result GateResult
	err    error
}

func (g fakeGate) Evaluate(_ context.Context, _ *Run, _ *Step, _ string) (GateResult, error) {
	return g.result, g.err
}

// ===== fakeAudit: orchestrator.AuditRecorder =====

type fakeAudit struct {
	mu      sync.Mutex
	actions []string
}

func (a *fakeAudit) Record(_ context.Context, _ uuid.UUID, action, _ string, _ map[string]any) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.actions = append(a.actions, action)
}
func (a *fakeAudit) has(action string) bool {
	a.mu.Lock()
	defer a.mu.Unlock()
	for _, x := range a.actions {
		if x == action {
			return true
		}
	}
	return false
}

// ===== fakeSessTerm: orchestrator.sessionTerminator =====

type fakeSessTerm struct {
	mu         sync.Mutex
	terminated []uuid.UUID
}

func (s *fakeSessTerm) TerminateSession(_ context.Context, _ uuid.UUID, _ bool, sessionID uuid.UUID) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.terminated = append(s.terminated, sessionID)
	return nil
}

// ===== fakeACL: orchestrator.ProjectACL =====

type fakeACL struct {
	allow bool
	read  bool
}

func (a fakeACL) CanRead(_ context.Context, _ uuid.UUID, _ bool, _ uuid.UUID) (bool, error) {
	return a.allow || a.read, nil
}

func (a fakeACL) CanWrite(_ context.Context, _ uuid.UUID, _ bool, _ uuid.UUID) (bool, error) {
	return a.allow, nil
}