fix: get tests passing again

coder · DanielleMaywood · Sep 24, 2024 · Sep 17, 2024 · Sep 17, 2024 · Sep 17, 2024
commit 6d829f3cd20083e75ddf1199366f98b258e303b8
diff --git a/agent/agent.go b/agent/agent.go
@@ -582,11 +582,6 @@ func (a *agent) reportMetadata(ctx context.Context, conn drpc.Conn) error {
 		select {
 		case <-ctx.Done():
 			return ctx.Err()
-		case timing := <-*a.scriptRunner.ScriptTimings():
-			_, err := aAPI.ScriptCompleted(ctx, &proto.WorkspaceAgentScriptCompletedRequest{
-				Timing: timing.ToProto(),
-			})
-			return err
 		case mr := <-metadataResults:
 			// This can overwrite unsent values, but that's fine because
 			// we're only interested about up-to-date values.
@@ -946,7 +941,7 @@ func (a *agent) handleManifest(manifestOK *checkpoint) func(ctx context.Context,
 				}
 			}
 
-			err = a.scriptRunner.Init(manifest.Scripts)
+			err = a.scriptRunner.Init(manifest.Scripts, aAPI.ScriptCompleted)
 			if err != nil {
 				return xerrors.Errorf("init script runner: %w", err)
 			}

diff --git a/agent/agentscripts/agentscripts.go b/agent/agentscripts/agentscripts.go
@@ -19,10 +19,12 @@ import (
 	"github.com/spf13/afero"
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/xerrors"
+	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"cdr.dev/slog"
 
 	"github.com/coder/coder/v2/agent/agentssh"
+	"github.com/coder/coder/v2/agent/proto"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 )
@@ -66,7 +68,6 @@ func New(opts Options) *Runner {
 		cronCtxCancel: cronCtxCancel,
 		cron:          cron.New(cron.WithParser(parser)),
 		closed:        make(chan struct{}),
-		scriptTimings: make(chan TimingSpan),
 		dataDir:       filepath.Join(opts.DataDirBase, "coder-script-data"),
 		scriptsExecuted: prometheus.NewCounterVec(prometheus.CounterOpts{
 			Namespace: "agent",
@@ -76,30 +77,28 @@ func New(opts Options) *Runner {
 	}
 }
 
+type ScriptCompletedFunc func(context.Context, *proto.WorkspaceAgentScriptCompletedRequest) (*proto.WorkspaceAgentScriptCompletedResponse, error)
+
 type Runner struct {
 	Options
 
-	cronCtx       context.Context
-	cronCtxCancel context.CancelFunc
-	cmdCloseWait  sync.WaitGroup
-	closed        chan struct{}
-	closeMutex    sync.Mutex
-	cron          *cron.Cron
-	initialized   atomic.Bool
-	scripts       []codersdk.WorkspaceAgentScript
-	scriptTimings chan TimingSpan
-	dataDir       string
+	cronCtx         context.Context
+	cronCtxCancel   context.CancelFunc
+	cmdCloseWait    sync.WaitGroup
+	closed          chan struct{}
+	closeMutex      sync.Mutex
+	cron            *cron.Cron
+	initialized     atomic.Bool
+	scripts         []codersdk.WorkspaceAgentScript
+	dataDir         string
+	scriptCompleted ScriptCompletedFunc
 
 	// scriptsExecuted includes all scripts executed by the workspace agent. Agents
 	// execute startup scripts, and scripts on a cron schedule. Both will increment
 	// this counter.
 	scriptsExecuted *prometheus.CounterVec
 }
 
-func (r *Runner) ScriptTimings() *chan TimingSpan {
-	return &r.scriptTimings
-}
-
 // DataDir returns the directory where scripts data is stored.
 func (r *Runner) DataDir() string {
 	return r.dataDir
@@ -122,12 +121,13 @@ func (r *Runner) RegisterMetrics(reg prometheus.Registerer) {
 // Init initializes the runner with the provided scripts.
 // It also schedules any scripts that have a schedule.
 // This function must be called before Execute.
-func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript) error {
+func (r *Runner) Init(scripts []codersdk.WorkspaceAgentScript, scriptCompleted ScriptCompletedFunc) error {
 	if r.initialized.Load() {
 		return xerrors.New("init: already initialized")
 	}
 	r.initialized.Store(true)
 	r.scripts = scripts
+	r.scriptCompleted = scriptCompleted
 	r.Logger.Info(r.cronCtx, "initializing agent scripts", slog.F("script_count", len(scripts)), slog.F("log_dir", r.LogDir))
 
 	err := r.Filesystem.MkdirAll(r.ScriptBinDir(), 0o700)
@@ -321,12 +321,14 @@ func (r *Runner) run(ctx context.Context, script codersdk.WorkspaceAgentScript)
 			logger.Info(ctx, fmt.Sprintf("%s script completed", logPath), slog.F("execution_time", execTime), slog.F("exit_code", exitCode))
 		}
 
-		r.scriptTimings <- TimingSpan{
-			displayName: script.DisplayName,
-			start:       start,
-			end:         end,
-			exitCode:    int32(exitCode),
-		}
+		_, err = r.scriptCompleted(ctx, &proto.WorkspaceAgentScriptCompletedRequest{
+			Timing: &proto.Timing{
+				DisplayName: script.DisplayName,
+				Start:       timestamppb.New(start),
+				End:         timestamppb.New(end),
+				ExitCode:    int32(exitCode),
+			},
+		})
 	}()
 
 	err = cmd.Start()

diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go
@@ -17,6 +17,7 @@ import (
 	"cdr.dev/slog/sloggers/slogtest"
 	"github.com/coder/coder/v2/agent/agentscripts"
 	"github.com/coder/coder/v2/agent/agentssh"
+	"github.com/coder/coder/v2/agent/agenttest"
 	"github.com/coder/coder/v2/codersdk"
 	"github.com/coder/coder/v2/codersdk/agentsdk"
 	"github.com/coder/coder/v2/testutil"
@@ -34,10 +35,11 @@ func TestExecuteBasic(t *testing.T) {
 		return fLogger
 	})
 	defer runner.Close()
+	aAPI := agenttest.NewFakeAgentAPI(t, slogtest.Make(t, nil), nil, nil)
 	err := runner.Init([]codersdk.WorkspaceAgentScript{{
 		LogSourceID: uuid.New(),
 		Script:      "echo hello",
-	}})
+	}}, aAPI.ScriptCompleted)
 	require.NoError(t, err)
 	require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
 		return true
@@ -61,10 +63,11 @@ func TestEnv(t *testing.T) {
 			cmd.exe /c echo %CODER_SCRIPT_BIN_DIR%
 		`
 	}
+	aAPI := agenttest.NewFakeAgentAPI(t, slogtest.Make(t, nil), nil, nil)
 	err := runner.Init([]codersdk.WorkspaceAgentScript{{
 		LogSourceID: id,
 		Script:      script,
-	}})
+	}}, aAPI.ScriptCompleted)
 	require.NoError(t, err)
 
 	ctx := testutil.Context(t, testutil.WaitLong)
@@ -103,11 +106,12 @@ func TestTimeout(t *testing.T) {
 	t.Parallel()
 	runner := setup(t, nil)
 	defer runner.Close()
+	aAPI := agenttest.NewFakeAgentAPI(t, slogtest.Make(t, nil), nil, nil)
 	err := runner.Init([]codersdk.WorkspaceAgentScript{{
 		LogSourceID: uuid.New(),
 		Script:      "sleep infinity",
 		Timeout:     time.Millisecond,
-	}})
+	}}, aAPI.ScriptCompleted)
 	require.NoError(t, err)
 	require.ErrorIs(t, runner.Execute(context.Background(), nil), agentscripts.ErrTimeout)
 }

diff --git a/agent/agentscripts/timings.go b/agent/agentscripts/timings.go