feat: self-update of master and agents (#15) (#32)

internal/e2e/jobs_test.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"log/slog"
 	"net/http"
+	"strconv"
 	"strings"
 	"testing"
 	"time"
@@ -46,7 +47,7 @@ func newGateExecutor() *gateExecutor {
 	return &gateExecutor{started: make(chan string, 8), release: make(chan struct{})}
 }
 
-func (g *gateExecutor) Execute(ctx context.Context, job protocol.JobStart, out io.Writer) (string, error) {
+func (g *gateExecutor) Execute(ctx context.Context, job updater.Job, out io.Writer) (string, error) {
 	_, _ = io.WriteString(out, "pulling "+job.Ref+"\n")
 	g.started <- job.Ref
 	select {
@@ -69,11 +70,17 @@ type jobAgent struct {
 	stopRun  func()
 }
 
-func startJobAgent(t *testing.T, m testMaster, token string, exec updater.Executor) *jobAgent {
+// startJobAgent starts an agent. before, if given, runs on the new runner
+// before the worker and the connection start, like ResumeSelfUpdates in the
+// real agent.
+func startJobAgent(t *testing.T, m testMaster, token string, exec updater.Executor, before ...func(*updater.Runner)) *jobAgent {
 	t.Helper()
 	src := &staticSource{}
 	src.set(web)
 	runner := updater.NewRunner(exec, time.Minute)
+	for _, f := range before {
+		f(runner)
+	}
 	runCtx, stopRun := context.WithCancel(context.Background())
 	runDone := make(chan struct{})
 	go func() { runner.Run(runCtx); close(runDone) }()
@@ -221,6 +228,54 @@ func TestAgentRestartMarksJobLost(t *testing.T) {
 	assert.Equal(t, lostMessage, lost.Error)
 }
 
+// TestAgentSelfUpdateResumedByNewAgent follows a self-update: the old agent
+// dies while its job runs, and the new agent finishes the job under the same
+// ID before it says hello, so the master neither loses nor re-runs it.
+func TestAgentSelfUpdateResumedByNewAgent(t *testing.T) {
+	pinned := slusk + ":1.4.2@" + e2eDigest
+	tests := []struct {
+		name      string
+		resume    updater.ResumeFunc
+		wantError string
+	}{
+		{"helper succeeded", func(_ context.Context, out io.Writer) (string, error) {
+			_, _ = io.WriteString(out, "helper output\n")
+			return pinned, nil
+		}, ""},
+		{"helper failed", func(context.Context, io.Writer) (string, error) {
+			return "", errors.New("docker compose up: exit status 1 (output: boom)")
+		}, "docker compose up: exit status 1 (output: boom)"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			e := newJobEnv(t)
+			j := e.trigger(t)
+			e.waitStarted(t)
+			started := time.Now().Add(-time.Second)
+
+			e.agent.stopConn()
+			e.agent.stopRun()
+			startJobAgent(t, e.m, "pua_ok", newGateExecutor(), func(r *updater.Runner) {
+				r.Resume(strconv.FormatInt(j.ID, 10), started, tt.resume)
+			})
+
+			if tt.wantError != "" {
+				done := e.waitStatus(t, j.ID, model.JobFailed)
+				assert.Equal(t, tt.wantError, done.Error)
+				assert.NotEqual(t, lostMessage, done.Error)
+				return
+			}
+			done := e.waitStatus(t, j.ID, model.JobSucceeded)
+			assert.Empty(t, done.Error)
+			assert.Equal(t, pinned, done.ToRef)
+			assert.Contains(t, done.Output, "helper output")
+			svc, err := e.store.GetService(context.Background(), 1)
+			require.NoError(t, err)
+			assert.Equal(t, pinned, svc.DesiredRef)
+		})
+	}
+}
+
 func TestPhaseOneAgentGetsNoJobs(t *testing.T) {
 	store := &memStore{agents: map[string]model.Agent{agenttoken.Hash("pua_old"): {ID: agentID, Name: "node-2"}}}
 	svcRow := model.Service{ID: 1, Kind: protocol.KindCompose, ComposeProject: "slusk", ComposeService: "web",