add agentproc tests

coder · sreya · Sep 15, 2023 · Aug 30, 2023 · Aug 31, 2023 · Aug 31, 2023
commit 8c652162f3c30ecbb71cae6607242b6b0388ab49
diff --git a/agent/agent.go b/agent/agent.go
@@ -1317,7 +1317,7 @@ func (a *agent) manageProcessPriorityLoop(ctx context.Context) {
 					continue
 				}
 
-				score, err := proc.Nice(a.syscaller)
+				score, err := proc.Niceness(a.syscaller)
 				if err != nil {
 					a.logger.Error(ctx, "unable to get proc niceness",
 						slog.F("name", proc.Name()),

diff --git a/agent/agentproc/agentproctest/proc.go b/agent/agentproc/agentproctest/proc.go
@@ -0,0 +1,44 @@
+package agentproctest
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/v2/agent/agentproc"
+	"github.com/coder/coder/v2/cryptorand"
+)
+
+func GenerateProcess(t *testing.T, fs afero.Fs, dir string) agentproc.Process {
+	t.Helper()
+
+	pid, err := cryptorand.Intn(1<<31 - 1)
+	require.NoError(t, err)
+
+	err = fs.MkdirAll(fmt.Sprintf("/%s/%d", dir, pid), 0555)
+	require.NoError(t, err)
+
+	arg1, err := cryptorand.String(5)
+	require.NoError(t, err)
+
+	arg2, err := cryptorand.String(5)
+	require.NoError(t, err)
+
+	arg3, err := cryptorand.String(5)
+	require.NoError(t, err)
+
+	cmdline := fmt.Sprintf("%s\x00%s\x00%s", arg1, arg2, arg3)
+
+	err = afero.WriteFile(fs, fmt.Sprintf("/%s/%d/cmdline", dir, pid), []byte(cmdline), 0444)
+	require.NoError(t, err)
+
+	return agentproc.Process{
+		PID:     int32(pid),
+		CmdLine: cmdline,
+		Dir:     fmt.Sprintf("%s/%d", dir, pid),
+		FS:      fs,
+	}
+
+}
diff --git a/agent/agentproc/doc.go b/agent/agentproc/doc.go
@@ -1,3 +1,5 @@
 // Package agentproc contains logic for interfacing with local
 // processes running in the same context as the agent.
 package agentproc
+
+//go:generate mockgen -destination ./syscallermock_test.go -package agentproc_test github.com/coder/coder/v2/agent/agentproc Syscaller
diff --git a/agent/agentproc/proc.go b/agent/agentproc/proc.go
@@ -1,61 +1,28 @@
 package agentproc
 
 import (
+	"errors"
 	"path/filepath"
 	"strconv"
 	"strings"
 	"syscall"
 
 	"github.com/spf13/afero"
-	"golang.org/x/sys/unix"
 	"golang.org/x/xerrors"
 )
 
 const DefaultProcDir = "/proc"
 
-type Syscaller interface {
-	SetPriority(pid int32, priority int) error
-	GetPriority(pid int32) (int, error)
-	Kill(pid int32, sig syscall.Signal) error
-}
-
-type UnixSyscaller struct{}
-
-func (UnixSyscaller) SetPriority(pid int32, nice int) error {
-	err := unix.Setpriority(unix.PRIO_PROCESS, int(pid), nice)
-	if err != nil {
-		return xerrors.Errorf("set priority: %w", err)
-	}
-	return nil
-}
-
-func (UnixSyscaller) GetPriority(pid int32) (int, error) {
-	nice, err := unix.Getpriority(0, int(pid))
-	if err != nil {
-		return 0, xerrors.Errorf("get priority: %w", err)
-	}
-	return nice, nil
-}
-
-func (UnixSyscaller) Kill(pid int, sig syscall.Signal) error {
-	err := syscall.Kill(pid, sig)
-	if err != nil {
-		return xerrors.Errorf("kill: %w", err)
-	}
-
-	return nil
-}
-
 type Process struct {
 	Dir     string
 	CmdLine string
 	PID     int32
-	fs      afero.Fs
+	FS      afero.Fs
 }
 
 func (p *Process) SetOOMAdj(score int) error {
 	path := filepath.Join(p.Dir, "oom_score_adj")
-	err := afero.WriteFile(p.fs,
+	err := afero.WriteFile(p.FS,
 		path,
 		[]byte(strconv.Itoa(score)),
 		0644,
@@ -67,20 +34,20 @@ func (p *Process) SetOOMAdj(score int) error {
 	return nil
 }
 
-func (p *Process) SetNiceness(sc Syscaller, score int) error {
-	err := sc.SetPriority(p.PID, score)
+func (p *Process) Niceness(sc Syscaller) (int, error) {
+	nice, err := sc.GetPriority(p.PID)
 	if err != nil {
-		return xerrors.Errorf("set priority for %q: %w", p.CmdLine, err)
+		return 0, xerrors.Errorf("get priority for %q: %w", p.CmdLine, err)
 	}
-	return nil
+	return nice, nil
 }
 
-func (p *Process) Nice(sc Syscaller) (int, error) {
-	nice, err := sc.GetPriority(p.PID)
+func (p *Process) SetNiceness(sc Syscaller, score int) error {
+	err := sc.SetPriority(p.PID, score)
 	if err != nil {
-		return 0, xerrors.Errorf("get priority for %q: %w", p.CmdLine, err)
+		return xerrors.Errorf("set priority for %q: %w", p.CmdLine, err)
 	}
-	return nice, nil
+	return nil
 }
 
 func (p *Process) Name() string {
@@ -108,7 +75,7 @@ func List(fs afero.Fs, syscaller Syscaller, dir string) ([]*Process, error) {
 		}
 
 		// Check that the process still exists.
-		exists, err := isProcessExist(syscaller, int32(pid), syscall.Signal(0))
+		exists, err := isProcessExist(syscaller, int32(pid))
 		if err != nil {
 			return nil, xerrors.Errorf("check process exists: %w", err)
 		}
@@ -128,26 +95,27 @@ func List(fs afero.Fs, syscaller Syscaller, dir string) ([]*Process, error) {
 			PID:     int32(pid),
 			CmdLine: string(cmdline),
 			Dir:     filepath.Join(dir, entry),
-			fs:      fs,
+			FS:      fs,
 		})
 	}
 
 	return processes, nil
 }
 
-func isProcessExist(syscaller Syscaller, pid int32, sig syscall.Signal) (bool, error) {
-	err := syscaller.Kill(pid, sig)
+func isProcessExist(syscaller Syscaller, pid int32) (bool, error) {
+	err := syscaller.Kill(pid, syscall.Signal(0))
 	if err == nil {
 		return true, nil
 	}
 	if err.Error() == "os: process already finished" {
 		return false, nil
 	}
 
-	errno, ok := err.(syscall.Errno)
-	if !ok {
+	var errno syscall.Errno
+	if !errors.As(err, &errno) {
 		return false, err
 	}
+
 	switch errno {
 	case syscall.ESRCH:
 		return false, nil

diff --git a/agent/agentproc/proc_test.go b/agent/agentproc/proc_test.go
@@ -1,12 +1,180 @@
 package agentproc_test
 
-type mockSyscaller struct {
-	SetPriorityFn func(int32, int) error
+import (
+	"fmt"
+	"strings"
+	"syscall"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/v2/agent/agentproc"
+	"github.com/coder/coder/v2/agent/agentproc/agentproctest"
+)
+
+func TestList(t *testing.T) {
+	t.Parallel()
+
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			fs            = afero.NewMemMapFs()
+			sc            = NewMockSyscaller(gomock.NewController(t))
+			expectedProcs = make(map[int32]agentproc.Process)
+			rootDir       = agentproc.DefaultProcDir
+		)
+
+		for i := 0; i < 4; i++ {
+			proc := agentproctest.GenerateProcess(t, fs, rootDir)
+			expectedProcs[proc.PID] = proc
+
+			sc.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+		}
+
+		actualProcs, err := agentproc.List(fs, sc, rootDir)
+		require.NoError(t, err)
+		require.Len(t, actualProcs, 4)
+		for _, proc := range actualProcs {
+			expected, ok := expectedProcs[proc.PID]
+			require.True(t, ok)
+			require.Equal(t, expected.PID, proc.PID)
+			require.Equal(t, expected.CmdLine, proc.CmdLine)
+			require.Equal(t, expected.Dir, proc.Dir)
+		}
+	})
+
+	t.Run("FinishedProcess", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			fs            = afero.NewMemMapFs()
+			sc            = NewMockSyscaller(gomock.NewController(t))
+			expectedProcs = make(map[int32]agentproc.Process)
+			rootDir       = agentproc.DefaultProcDir
+		)
+
+		for i := 0; i < 3; i++ {
+			proc := agentproctest.GenerateProcess(t, fs, rootDir)
+			expectedProcs[proc.PID] = proc
+
+			sc.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+		}
+
+		// Create a process that's already finished. We're not adding
+		// it to the map because it should be skipped over.
+		proc := agentproctest.GenerateProcess(t, fs, rootDir)
+		sc.EXPECT().
+			Kill(proc.PID, syscall.Signal(0)).
+			Return(xerrors.New("os: process already finished"))
+
+		actualProcs, err := agentproc.List(fs, sc, rootDir)
+		require.NoError(t, err)
+		require.Len(t, actualProcs, 3)
+		for _, proc := range actualProcs {
+			expected, ok := expectedProcs[proc.PID]
+			require.True(t, ok)
+			require.Equal(t, expected.PID, proc.PID)
+			require.Equal(t, expected.CmdLine, proc.CmdLine)
+			require.Equal(t, expected.Dir, proc.Dir)
+		}
+	})
+
+	t.Run("NoSuchProcess", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			fs            = afero.NewMemMapFs()
+			sc            = NewMockSyscaller(gomock.NewController(t))
+			expectedProcs = make(map[int32]agentproc.Process)
+			rootDir       = agentproc.DefaultProcDir
+		)
+
+		for i := 0; i < 3; i++ {
+			proc := agentproctest.GenerateProcess(t, fs, rootDir)
+			expectedProcs[proc.PID] = proc
+
+			sc.EXPECT().
+				Kill(proc.PID, syscall.Signal(0)).
+				Return(nil)
+		}
+
+		// Create a process that doesn't exist. We're not adding
+		// it to the map because it should be skipped over.
+		proc := agentproctest.GenerateProcess(t, fs, rootDir)
+		sc.EXPECT().
+			Kill(proc.PID, syscall.Signal(0)).
+			Return(syscall.ESRCH)
+
+		actualProcs, err := agentproc.List(fs, sc, rootDir)
+		require.NoError(t, err)
+		require.Len(t, actualProcs, 3)
+		for _, proc := range actualProcs {
+			expected, ok := expectedProcs[proc.PID]
+			require.True(t, ok)
+			require.Equal(t, expected.PID, proc.PID)
+			require.Equal(t, expected.CmdLine, proc.CmdLine)
+			require.Equal(t, expected.Dir, proc.Dir)
+		}
+	})
 }
 
-func (f mockSyscaller) SetPriority(pid int32, nice int) error {
-	if f.SetPriorityFn == nil {
-		return nil
-	}
-	return f.SetPriorityFn(pid, nice)
+// These tests are not very interesting but they provide some modicum of
+// confidence.
+func TestProcess(t *testing.T) {
+	t.Parallel()
+
+	t.Run("SetOOMAdj", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			fs            = afero.NewMemMapFs()
+			dir           = agentproc.DefaultProcDir
+			proc          = agentproctest.GenerateProcess(t, fs, agentproc.DefaultProcDir)
+			expectedScore = -1000
+		)
+
+		err := proc.SetOOMAdj(expectedScore)
+		require.NoError(t, err)
+
+		actualScore, err := afero.ReadFile(fs, fmt.Sprintf("%s/%d/oom_score_adj", dir, proc.PID))
+		require.NoError(t, err)
+		require.Equal(t, fmt.Sprintf("%d", expectedScore), strings.TrimSpace(string(actualScore)))
+	})
+
+	t.Run("SetNiceness", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			sc   = NewMockSyscaller(gomock.NewController(t))
+			proc = &agentproc.Process{
+				PID: 32,
+			}
+			score = 20
+		)
+
+		sc.EXPECT().SetPriority(proc.PID, score).Return(nil)
+		err := proc.SetNiceness(sc, score)
+		require.NoError(t, err)
+	})
+
+	t.Run("Name", func(t *testing.T) {
+		t.Parallel()
+
+		var (
+			proc = &agentproc.Process{
+				CmdLine: "helloworld\x00--arg1\x00--arg2",
+			}
+			expectedName = "helloworld"
+		)
+
+		require.Equal(t, expectedName, proc.Name())
+	})
 }