add agentscripts test for execute option

mafredri · mafredri · commit c78665ae1871 · 2025-03-26T13:37:49.000Z
diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"path/filepath"
 	"runtime"
+	"slices"
+	"sync"
 	"testing"
 	"time"
 
@@ -151,6 +153,154 @@ func TestCronClose(t *testing.T) {
 	require.NoError(t, runner.Close(), "close runner")
 }
 
+func TestExecuteOptions(t *testing.T) {
+	t.Parallel()
+
+	startScript := codersdk.WorkspaceAgentScript{
+		ID:          uuid.New(),
+		LogSourceID: uuid.New(),
+		Script:      "echo start",
+		RunOnStart:  true,
+	}
+	stopScript := codersdk.WorkspaceAgentScript{
+		ID:          uuid.New(),
+		LogSourceID: uuid.New(),
+		Script:      "echo stop",
+		RunOnStop:   true,
+	}
+	postStartScript := codersdk.WorkspaceAgentScript{
+		ID:          uuid.New(),
+		LogSourceID: uuid.New(),
+		Script:      "echo poststart",
+	}
+	regularScript := codersdk.WorkspaceAgentScript{
+		ID:          uuid.New(),
+		LogSourceID: uuid.New(),
+		Script:      "echo regular",
+	}
+
+	scripts := []codersdk.WorkspaceAgentScript{
+		startScript,
+		stopScript,
+		regularScript,
+	}
+	allScripts := append(slices.Clone(scripts), postStartScript)
+
+	scriptByID := func(id uuid.UUID) codersdk.WorkspaceAgentScript {
+		for _, script := range allScripts {
+			if script.ID == id {
+				return script
+			}
+		}
+		return codersdk.WorkspaceAgentScript{}
+	}
+
+	wantOutput := map[uuid.UUID]string{
+		startScript.ID:     "start",
+		stopScript.ID:      "stop",
+		postStartScript.ID: "poststart",
+		regularScript.ID:   "regular",
+	}
+
+	testCases := []struct {
+		name    string
+		option  agentscripts.ExecuteOption
+		wantRun []uuid.UUID
+	}{
+		{
+			name:    "ExecuteAllScripts",
+			option:  agentscripts.ExecuteAllScripts,
+			wantRun: []uuid.UUID{startScript.ID, stopScript.ID, regularScript.ID, postStartScript.ID},
+		},
+		{
+			name:    "ExecuteStartScripts",
+			option:  agentscripts.ExecuteStartScripts,
+			wantRun: []uuid.UUID{startScript.ID},
+		},
+		{
+			name:    "ExecutePostStartScripts",
+			option:  agentscripts.ExecutePostStartScripts,
+			wantRun: []uuid.UUID{postStartScript.ID},
+		},
+		{
+			name:    "ExecuteStopScripts",
+			option:  agentscripts.ExecuteStopScripts,
+			wantRun: []uuid.UUID{stopScript.ID},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ctx := testutil.Context(t, testutil.WaitMedium)
+			executedScripts := make(map[uuid.UUID]bool)
+			fLogger := &filterTestLogger{
+				tb:              t,
+				executedScripts: executedScripts,
+				wantOutput:      wantOutput,
+			}
+
+			runner := setup(t, func(_ uuid.UUID) agentscripts.ScriptLogger {
+				return fLogger
+			})
+			defer runner.Close()
+
+			aAPI := agenttest.NewFakeAgentAPI(t, testutil.Logger(t), nil, nil)
+			err := runner.Init(
+				scripts,
+				aAPI.ScriptCompleted,
+				agentscripts.WithPostStartScripts(postStartScript),
+			)
+			require.NoError(t, err)
+
+			err = runner.Execute(ctx, tc.option)
+			require.NoError(t, err)
+
+			gotRun := map[uuid.UUID]bool{}
+			for _, id := range tc.wantRun {
+				gotRun[id] = true
+				require.True(t, executedScripts[id],
+					"script %s should have run when using filter %s", scriptByID(id).Script, tc.name)
+			}
+
+			for _, script := range allScripts {
+				if _, ok := gotRun[script.ID]; ok {
+					continue
+				}
+				require.False(t, executedScripts[script.ID],
+					"script %s should not have run when using filter %s", script.Script, tc.name)
+			}
+		})
+	}
+}
+
+type filterTestLogger struct {
+	tb              testing.TB
+	executedScripts map[uuid.UUID]bool
+	wantOutput      map[uuid.UUID]string
+	mu              sync.Mutex
+}
+
+func (l *filterTestLogger) Send(ctx context.Context, logs ...agentsdk.Log) error {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	for _, log := range logs {
+		l.tb.Log(log.Output)
+		for id, output := range l.wantOutput {
+			if log.Output == output {
+				l.executedScripts[id] = true
+				break
+			}
+		}
+	}
+	return nil
+}
+
+func (l *filterTestLogger) Flush(context.Context) error {
+	return nil
+}
+
 func setup(t *testing.T, getScriptLogger func(logSourceID uuid.UUID) agentscripts.ScriptLogger) *agentscripts.Runner {
 	t.Helper()
 	if getScriptLogger == nil {
