coder/agent/agentscripts/agentscripts_test.go

package agentscripts_test

import (
	"context"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/spf13/afero"
	"github.com/stretchr/testify/require"
	"go.uber.org/atomic"
	"go.uber.org/goleak"

	"cdr.dev/slog/sloggers/slogtest"
	"github.com/coder/coder/v2/agent/agentscripts"
	"github.com/coder/coder/v2/agent/agentssh"
	"github.com/coder/coder/v2/codersdk"
	"github.com/coder/coder/v2/codersdk/agentsdk"
)

func TestMain(m *testing.M) {
	goleak.VerifyTestMain(m)
}

func TestExecuteBasic(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()
	err := runner.Init([]codersdk.WorkspaceAgentScript{{
		Script: "echo hello",
	}})
	require.NoError(t, err)
	require.NoError(t, runner.Execute(context.Background(), func(script codersdk.WorkspaceAgentScript) bool {
		return true
	}))
	log := <-logs
	require.Equal(t, "hello", log.Logs[0].Output)
}

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,
	}})
	require.NoError(t, err)
	require.ErrorIs(t, runner.Execute(context.Background(), nil), agentscripts.ErrTimeout)
}

func setup(t *testing.T, patchLogs func(ctx context.Context, req agentsdk.PatchLogs) error) *agentscripts.Runner {
	t.Helper()
	if patchLogs == nil {
		// noop
		patchLogs = func(ctx context.Context, req agentsdk.PatchLogs) error {
			return nil
		}
	}
	fs := afero.NewMemMapFs()
	logger := slogtest.Make(t, nil)
	s, err := agentssh.NewServer(context.Background(), logger, prometheus.NewRegistry(), fs, 0, "")
	require.NoError(t, err)
	s.AgentToken = func() string { return "" }
	s.Manifest = atomic.NewPointer(&agentsdk.Manifest{})
	t.Cleanup(func() {
		_ = s.Close()
	})
	return agentscripts.New(agentscripts.Options{
		LogDir:     t.TempDir(),
		Logger:     logger,
		SSHServer:  s,
		Filesystem: fs,
		PatchLogs:  patchLogs,
	})
}