add agentscript test

mafredri · mafredri · commit b0445eb07a8c · 2024-02-19T13:50:01.000+02:00
diff --git a/agent/agentscripts/agentscripts_test.go b/agent/agentscripts/agentscripts_test.go
@@ -2,9 +2,11 @@ package agentscripts_test
 
 import (
 	"context"
+	"path/filepath"
 	"testing"
 	"time"
 
+	"github.com/google/uuid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/spf13/afero"
 	"github.com/stretchr/testify/require"
@@ -30,7 +32,8 @@ func TestExecuteBasic(t *testing.T) {
 	})
 	defer runner.Close()
 	err := runner.Init([]codersdk.WorkspaceAgentScript{{
-		Script: "echo hello",
+		LogSourceID: uuid.New(),
+		Script:      "echo hello",
 	}})
 	require.NoError(t, err)
 	require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
@@ -40,13 +43,36 @@ func TestExecuteBasic(t *testing.T) {
 	require.Equal(t, "hello", log.Logs[0].Output)
 }
 
+func TestEnv(t *testing.T) {
+	t.Parallel()
+	logs := make(chan agentsdk.PatchLogs, 1)
+	runner := setup(t, func(ctx context.Context, req agentsdk.PatchLogs) error {
+		logs <- req
+		return nil
+	})
+	defer runner.Close()
+	id := uuid.New()
+	err := runner.Init([]codersdk.WorkspaceAgentScript{{
+		LogSourceID: id,
+		Script:      "echo $CODER_SCRIPT_DATA_DIR\necho $CODER_SCRIPT_BIN_DIR\n",
+	}})
+	require.NoError(t, err)
+	require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
+		return true
+	}))
+	log := <-logs
+	require.Contains(t, log.Logs[0].Output, filepath.Join(runner.DataDir, "coder-script-data", id.String()))
+	require.Contains(t, log.Logs[1].Output, runner.ScriptBinDir())
+}
+
 func TestTimeout(t *testing.T) {
 	t.Parallel()
 	runner := setup(t, nil)
 	defer runner.Close()
 	err := runner.Init([]codersdk.WorkspaceAgentScript{{
-		Script:  "sleep infinity",
-		Timeout: time.Millisecond,
+		LogSourceID: uuid.New(),
+		Script:      "sleep infinity",
+		Timeout:     time.Millisecond,
 	}})
 	require.NoError(t, err)
 	require.ErrorIs(t, runner.Execute(context.Background(), nil), agentscripts.ErrTimeout)
@@ -78,6 +104,7 @@ func setup(t *testing.T, patchLogs func(ctx context.Context, req agentsdk.PatchL
 	})
 	return agentscripts.New(agentscripts.Options{
 		LogDir:     t.TempDir(),
+		DataDir:    t.TempDir(),
 		Logger:     logger,
 		SSHServer:  s,
 		Filesystem: fs,
