package-updater/internal/agent/updater/runner_test.go

package updater

import (
	"context"
	"errors"
	"io"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/shcizo/package-updater/internal/protocol"
)

// fakeExec writes "pulling <ref>" and then waits on gate (if set). Refs
// starting with "bad" fail. It records the highest concurrency it saw.
type fakeExec struct {
	gate    chan struct{}
	cur     atomic.Int32
	maxSeen atomic.Int32
	started chan string
}

func newFakeExec() *fakeExec {
	return &fakeExec{gate: make(chan struct{}), started: make(chan string, 10)}
}

func (f *fakeExec) Execute(ctx context.Context, job Job, out io.Writer) (string, error) {
	n := f.cur.Add(1)
	defer f.cur.Add(-1)
	for {
		m := f.maxSeen.Load()
		if n <= m || f.maxSeen.CompareAndSwap(m, n) {
			break
		}
	}
	f.started <- job.Ref
	io.WriteString(out, "pulling "+job.Ref+"\n")
	select {
	case <-f.gate:
	case <-ctx.Done():
		return "", ctx.Err()
	}
	if strings.HasPrefix(job.Ref, "bad") {
		return "", errors.New("boom")
	}
	return job.Ref + "@sha256:aa", nil
}

func startRunner(t *testing.T, exec Executor, timeout time.Duration) *Runner {
	t.Helper()
	r := NewRunner(exec, timeout)
	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); r.Run(ctx) }()
	t.Cleanup(func() { cancel(); wg.Wait() })
	return r
}

func waitPending(t *testing.T, r *Runner, n int) []Finished {
	t.Helper()
	var got []Finished
	require.Eventually(t, func() bool { got = r.Pending(); return len(got) >= n }, 2*time.Second, 5*time.Millisecond)
	return got
}

func TestRunner_FIFOSingleWorker(t *testing.T) {
	fe := newFakeExec()
	r := startRunner(t, fe, time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	r.Submit("2", protocol.JobStart{Ref: "bad:2"})
	assert.Equal(t, "a:1", <-fe.started)
	fe.gate <- struct{}{}
	assert.Equal(t, "bad:2", <-fe.started)
	fe.gate <- struct{}{}

	got := waitPending(t, r, 2)
	assert.Equal(t, int32(1), fe.maxSeen.Load())
	require.Len(t, got, 2)
	assert.Equal(t, "1", got[0].ID)
	assert.Equal(t, protocol.ResultSucceeded, got[0].Result.Status)
	assert.Equal(t, "a:1@sha256:aa", got[0].Result.Ref)
	assert.Equal(t, "pulling a:1\n", got[0].Result.Output)
	assert.Equal(t, "2", got[1].ID)
	assert.Equal(t, protocol.ResultFailed, got[1].Result.Status)
	assert.Equal(t, "boom", got[1].Result.Error)
	assert.False(t, got[0].Result.StartedAt.IsZero())
}

func TestRunner_DuplicateSubmitIgnored(t *testing.T) {
	fe := newFakeExec()
	r := startRunner(t, fe, time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	assert.Equal(t, "a:1", <-fe.started)
	r.Submit("1", protocol.JobStart{Ref: "other:9"})
	fe.gate <- struct{}{}
	got := waitPending(t, r, 1)
	assert.Equal(t, "a:1@sha256:aa", got[0].Result.Ref)
	select {
	case ref := <-fe.started:
		t.Fatalf("duplicate job ran: %s", ref)
	case <-time.After(50 * time.Millisecond):
	}
}

func TestRunner_KnownAndAck(t *testing.T) {
	fe := newFakeExec()
	r := startRunner(t, fe, time.Minute)
	r.Submit("2", protocol.JobStart{Ref: "a:2"})
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	assert.Equal(t, "a:2", <-fe.started)
	assert.Equal(t, []string{"1", "2"}, r.Known())

	r.Ack("2") // running: ignored
	assert.Equal(t, []string{"1", "2"}, r.Known())

	fe.gate <- struct{}{}
	waitPending(t, r, 1)
	r.Ack("2")
	assert.Equal(t, []string{"1"}, r.Known())
	<-fe.started
	fe.gate <- struct{}{}
	waitPending(t, r, 1)
}

func TestRunner_TakeProgress(t *testing.T) {
	fe := newFakeExec()
	r := startRunner(t, fe, time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	<-fe.started
	require.Eventually(t, func() bool { return len(r.Pending()) == 0 && len(r.Known()) == 1 }, time.Second, 5*time.Millisecond)

	var first []Progress
	require.Eventually(t, func() bool { first = r.TakeProgress(); return len(first) > 0 }, time.Second, 5*time.Millisecond)
	assert.Equal(t, []Progress{{ID: "1", Output: "pulling a:1\n"}}, first)
	assert.Empty(t, r.TakeProgress(), "nothing new since the last call")

	fe.gate <- struct{}{}
	waitPending(t, r, 1)
	assert.Empty(t, r.TakeProgress(), "a finished job no longer reports progress")
}

func TestRunner_TakeProgressReportsStartedJobOnce(t *testing.T) {
	gate := make(chan struct{})
	started := make(chan struct{})
	r := startRunner(t, execFunc(func(ctx context.Context, _ Job, _ io.Writer) (string, error) {
		close(started)
		<-gate
		return "x", nil
	}), time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	<-started
	assert.Equal(t, []Progress{{ID: "1", Output: ""}}, r.TakeProgress())
	assert.Empty(t, r.TakeProgress())
	close(gate)
}

type execFunc func(ctx context.Context, job Job, out io.Writer) (string, error)

func (f execFunc) Execute(ctx context.Context, job Job, out io.Writer) (string, error) {
	return f(ctx, job, out)
}

func TestRunner_Timeout(t *testing.T) {
	r := startRunner(t, execFunc(func(ctx context.Context, _ Job, _ io.Writer) (string, error) {
		<-ctx.Done()
		return "", ctx.Err()
	}), 20*time.Millisecond)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	got := waitPending(t, r, 1)
	assert.Equal(t, protocol.ResultFailed, got[0].Result.Status)
	assert.Equal(t, "timed out after 20ms", got[0].Result.Error)
}

func TestRunner_OutputCapped(t *testing.T) {
	r := startRunner(t, execFunc(func(_ context.Context, _ Job, out io.Writer) (string, error) {
		io.WriteString(out, strings.Repeat("a", 10))
		io.WriteString(out, strings.Repeat("b", protocol.MaxJobOutput))
		return "x", nil
	}), time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	got := waitPending(t, r, 1)
	assert.Equal(t, strings.Repeat("b", protocol.MaxJobOutput), got[0].Result.Output)
	assert.True(t, got[0].Result.OutputTruncated)
}

func TestRunner_ChangedSignals(t *testing.T) {
	r := startRunner(t, execFunc(func(context.Context, Job, io.Writer) (string, error) { return "x", nil }), time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	select {
	case <-r.Changed():
	case <-time.After(time.Second):
		t.Fatal("no change signal")
	}
}

func TestRunner_ResumeJumpsTheQueue(t *testing.T) {
	fe := newFakeExec()
	r := NewRunner(fe, time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	started := time.Date(2026, 10, 13, 9, 0, 0, 0, time.UTC)
	r.Resume("9", started, func(_ context.Context, out io.Writer) (string, error) {
		io.WriteString(out, "from the helper\n")
		return "resumed@sha256:bb", nil
	})
	assert.Equal(t, []string{"1", "9"}, r.Known(), "known before the worker starts")

	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); r.Run(ctx) }()
	t.Cleanup(func() { cancel(); wg.Wait() })

	got := waitPending(t, r, 1)
	require.Len(t, got, 1, "the resumed job runs first; job 1 waits on the gate")
	assert.Equal(t, "9", got[0].ID)
	assert.Equal(t, protocol.ResultSucceeded, got[0].Result.Status)
	assert.Equal(t, "resumed@sha256:bb", got[0].Result.Ref)
	assert.Equal(t, "from the helper\n", got[0].Result.Output)
	assert.True(t, started.Equal(got[0].Result.StartedAt))

	r.Resume("9", started, nil) // already held: ignored
	assert.Equal(t, "a:1", <-fe.started)
	fe.gate <- struct{}{}
	waitPending(t, r, 2)
}

func TestRunner_ResumedJobFailure(t *testing.T) {
	r := startRunner(t, newFakeExec(), time.Minute)
	r.Resume("9", time.Now(), func(context.Context, io.Writer) (string, error) { return "", errors.New("helper failed") })
	got := waitPending(t, r, 1)
	assert.Equal(t, protocol.ResultFailed, got[0].Result.Status)
	assert.Equal(t, "helper failed", got[0].Result.Error)
}

func TestRunner_SelfJobsRunLast(t *testing.T) {
	fe := newFakeExec()
	r := NewRunner(fe, time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "self:1", Self: true})
	r.Submit("2", protocol.JobStart{Ref: "a:2"})
	r.Submit("3", protocol.JobStart{Ref: "self:3", Self: true})
	r.Submit("4", protocol.JobStart{Ref: "a:4"})

	ctx, cancel := context.WithCancel(context.Background())
	var wg sync.WaitGroup
	wg.Add(1)
	go func() { defer wg.Done(); r.Run(ctx) }()
	t.Cleanup(func() { cancel(); wg.Wait() })

	var order []string
	for i := 0; i < 4; i++ {
		order = append(order, <-fe.started)
		fe.gate <- struct{}{}
	}
	assert.Equal(t, []string{"a:2", "a:4", "self:1", "self:3"}, order)
}

func TestRunner_SelfJobRunsWhenAlone(t *testing.T) {
	fe := newFakeExec()
	r := startRunner(t, fe, time.Minute)
	r.Submit("1", protocol.JobStart{Ref: "self:1", Self: true})
	assert.Equal(t, "self:1", <-fe.started)
	fe.gate <- struct{}{}
	waitPending(t, r, 1)
}

func TestRunner_StoppingTheAgentRecordsNoResult(t *testing.T) {
	r := NewRunner(execFunc(func(ctx context.Context, _ Job, _ io.Writer) (string, error) {
		<-ctx.Done()
		return "", ctx.Err()
	}), time.Minute)
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan struct{})
	go func() { r.Run(ctx); close(done) }()
	r.Submit("1", protocol.JobStart{Ref: "a:1"})
	require.Eventually(t, func() bool { return len(r.TakeProgress()) > 0 }, time.Second, 5*time.Millisecond)

	cancel()
	<-done
	assert.Empty(t, r.Pending(), "no \"context canceled\" result")
	assert.Empty(t, r.Known(), "the job is dropped, so the master marks it lost unless a successor resumes it")
}

func TestRunner_ExecutorGetsIDAndStart(t *testing.T) {
	var got Job
	r := startRunner(t, execFunc(func(_ context.Context, job Job, _ io.Writer) (string, error) {
		got = job
		return "x", nil
	}), time.Minute)
	now := time.Date(2026, 10, 13, 9, 0, 0, 0, time.UTC)
	r.now = func() time.Time { return now }
	r.Submit("5", protocol.JobStart{Ref: "a:1"})
	waitPending(t, r, 1)
	assert.Equal(t, "5", got.ID)
	assert.Equal(t, "a:1", got.Ref)
	assert.True(t, now.Equal(got.Started))
}