Update unit tests for env collectors

stateful · Oct 13, 2024 · d9d4b71 · d9d4b71
1 parent 7351fbf
commit d9d4b71
Show file tree

Hide file tree

Showing 7 changed files with 95 additions and 36 deletions.
diff --git a/internal/cmd/beta/session_cmd.go b/internal/cmd/beta/session_cmd.go
@@ -109,7 +109,7 @@ func sessionSetupCmd() *cobra.Command {
 						return nil
 					}
 
-					envSetter := command.NewFileBasedEnvSetter(
+					envSetter := command.NewScriptEnvSetter(
 						os.Getenv(command.EnvCollectorSessionPrePathEnvName),
 						os.Getenv(command.EnvCollectorSessionPostPathEnvName),
 						debug,

diff --git a/internal/command/env_collector_fifo_unix.go b/internal/command/env_collector_fifo_unix.go
@@ -113,15 +113,18 @@ func (c *envCollectorFifo) Diff() (changed []string, deleted []string, _ error)
 }
 
 func (c *envCollectorFifo) ExtraEnv() []string {
-	if c.encKey == nil || c.encNonce == nil {
-		return nil
-	}
-	return []string{
-		createEnv(envCollectorEncKeyEnvName, hex.EncodeToString(c.encKey)),
-		createEnv(envCollectorEncNonceEnvName, hex.EncodeToString(c.encNonce)),
+	result := []string{
 		createEnv(EnvCollectorSessionPrePathEnvName, c.prePath()),
 		createEnv(EnvCollectorSessionPostPathEnvName, c.postPath()),
 	}
+	if c.encKey != nil && c.encNonce != nil {
+		result = append(
+			result,
+			createEnv(envCollectorEncKeyEnvName, hex.EncodeToString(c.encKey)),
+			createEnv(envCollectorEncNonceEnvName, hex.EncodeToString(c.encNonce)),
+		)
+	}
+	return result
 }
 
 func (c *envCollectorFifo) SetOnShell(shell io.Writer) error {

diff --git a/internal/command/env_collector_fifo_unix_test.go b/internal/command/env_collector_fifo_unix_test.go
@@ -4,13 +4,14 @@
 package command
 
 import (
+	"bytes"
 	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
-func TestEnvCollectorFifo(t *testing.T) {
+func Test_envCollectorFifo(t *testing.T) {
 	t.Parallel()
 
 	collector, err := newEnvCollectorFifo(scanEnv, nil, nil)
@@ -21,13 +22,29 @@ func TestEnvCollectorFifo(t *testing.T) {
 	err = os.WriteFile(collector.postPath(), []byte("ENV_2=2"), 0o600)
 	require.NoError(t, err)
 
-	changedEnv, deletedEnv, err := collector.Diff()
-	require.NoError(t, err)
-	require.Equal(t, []string{"ENV_2=2"}, changedEnv)
-	require.Equal(t, []string{"ENV_1"}, deletedEnv)
+	t.Run("ExtraEnv", func(t *testing.T) {
+		require.Len(t, collector.ExtraEnv(), 2)
+	})
+
+	t.Run("SetOnShell", func(t *testing.T) {
+		buf := new(bytes.Buffer)
+		err := collector.SetOnShell(buf)
+		require.NoError(t, err)
+		expected := " env -0 > " + collector.prePath() + "\n" +
+			" __cleanup() {\nrv=$?\nenv -0 > " + collector.postPath() + "\nexit $rv\n}\n" +
+			" trap -- \"__cleanup\" EXIT\n"
+		require.Equal(t, expected, buf.String())
+	})
+
+	t.Run("Diff", func(t *testing.T) {
+		changedEnv, deletedEnv, err := collector.Diff()
+		require.NoError(t, err)
+		require.Equal(t, []string{"ENV_2=2"}, changedEnv)
+		require.Equal(t, []string{"ENV_1"}, deletedEnv)
+	})
 }
 
-func TestEnvCollectorFifoWithoutWriter(t *testing.T) {
+func Test_envCollectorFifo_WithoutWriter(t *testing.T) {
 	t.Parallel()
 
 	collector, err := newEnvCollectorFifo(scanEnv, nil, nil)

diff --git a/internal/command/env_collector_file.go b/internal/command/env_collector_file.go
@@ -61,15 +61,18 @@ func (c *envCollectorFile) Diff() (changed []string, deleted []string, _ error)
 }
 
 func (c *envCollectorFile) ExtraEnv() []string {
-	if c.encKey == nil || c.encNonce == nil {
-		return nil
-	}
-	return []string{
-		createEnv(envCollectorEncKeyEnvName, hex.EncodeToString(c.encKey)),
-		createEnv(envCollectorEncNonceEnvName, hex.EncodeToString(c.encNonce)),
+	result := []string{
 		createEnv(EnvCollectorSessionPrePathEnvName, c.prePath()),
 		createEnv(EnvCollectorSessionPostPathEnvName, c.postPath()),
 	}
+	if c.encKey != nil && c.encNonce != nil {
+		result = append(
+			result,
+			createEnv(envCollectorEncKeyEnvName, hex.EncodeToString(c.encKey)),
+			createEnv(envCollectorEncNonceEnvName, hex.EncodeToString(c.encNonce)),
+		)
+	}
+	return result
 }
 
 func (c *envCollectorFile) SetOnShell(shell io.Writer) error {

diff --git a/internal/command/env_collector_file_test.go b/internal/command/env_collector_file_test.go
@@ -1,13 +1,14 @@
 package command
 
 import (
+	"bytes"
 	"os"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
-func TestEnvCollectorFile(t *testing.T) {
+func Test_envCollectorFile(t *testing.T) {
 	t.Parallel()
 
 	collector, err := newEnvCollectorFile(scanEnv, nil, nil)
@@ -18,8 +19,24 @@ func TestEnvCollectorFile(t *testing.T) {
 	err = os.WriteFile(collector.postPath(), []byte("ENV_2=2"), 0o600)
 	require.NoError(t, err)
 
-	changedEnv, deletedEnv, err := collector.Diff()
-	require.NoError(t, err)
-	require.Equal(t, []string{"ENV_2=2"}, changedEnv)
-	require.Equal(t, []string{"ENV_1"}, deletedEnv)
+	t.Run("ExtraEnv", func(t *testing.T) {
+		require.Len(t, collector.ExtraEnv(), 2)
+	})
+
+	t.Run("SetOnShell", func(t *testing.T) {
+		buf := new(bytes.Buffer)
+		err := collector.SetOnShell(buf)
+		require.NoError(t, err)
+		expected := " env -0 > " + collector.prePath() + "\n" +
+			" __cleanup() {\nrv=$?\nenv -0 > " + collector.postPath() + "\nexit $rv\n}\n" +
+			" trap -- \"__cleanup\" EXIT\n"
+		require.Equal(t, expected, buf.String())
+	})
+
+	t.Run("Diff", func(t *testing.T) {
+		changedEnv, deletedEnv, err := collector.Diff()
+		require.NoError(t, err)
+		require.Equal(t, []string{"ENV_2=2"}, changedEnv)
+		require.Equal(t, []string{"ENV_1"}, deletedEnv)
+	})
 }
diff --git a/internal/command/env_shell.go b/internal/command/env_shell.go
@@ -23,7 +23,7 @@ type ScriptEnvSetter struct {
 	postPath    string
 }
 
-func NewFileBasedEnvSetter(prePath, postPath string, debug bool) *ScriptEnvSetter {
+func NewScriptEnvSetter(prePath, postPath string, debug bool) *ScriptEnvSetter {
 	return &ScriptEnvSetter{
 		debug:       debug,
 		dumpCommand: envDumpCommand,

diff --git a/internal/command/env_shell_test.go b/internal/command/env_shell_test.go
@@ -10,22 +10,41 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestSetOnShell(t *testing.T) {
+func TestScriptEnvSetter(t *testing.T) {
 	t.Parallel()
 
-	buf := new(bytes.Buffer)
+	prePath := "/tmp/pre-path"
+	postPath := "/tmp/post-path"
 
-	err := setOnShell(buf, envDumpCommand, false, false, false, "prePath", "postPath")
-	require.NoError(t, err)
+	t.Run("WithDebug", func(t *testing.T) {
+		setter := NewScriptEnvSetter(prePath, postPath, true)
+		buf := new(bytes.Buffer)
 
-	expected := (envDumpCommand + " > prePath\n" +
-		"__cleanup() {\n" +
-		"rv=$?\n" +
-		envDumpCommand + " > postPath\n" +
-		"exit $rv\n}\n" +
-		"trap -- \"__cleanup\" EXIT\n")
+		err := setter.SetOnShell(buf)
+		require.NoError(t, err)
 
-	require.EqualValues(t, expected, buf.String())
+		expected := "#!/bin/sh\n" +
+			"set -euxo pipefail\n" +
+			"env -0 > /tmp/pre-path\n" +
+			"__cleanup() {\nrv=$?\nenv -0 > /tmp/post-path\nexit $rv\n}\n" +
+			"trap -- \"__cleanup\" EXIT\n" +
+			"set +euxo pipefail\n"
+		require.EqualValues(t, expected, buf.String())
+	})
+
+	t.Run("WithoutDebug", func(t *testing.T) {
+		setter := NewScriptEnvSetter(prePath, postPath, false)
+		buf := new(bytes.Buffer)
+
+		err := setter.SetOnShell(buf)
+		require.NoError(t, err)
+
+		expected := "#!/bin/sh\n" +
+			"env -0 > /tmp/pre-path\n" +
+			"__cleanup() {\nrv=$?\nenv -0 > /tmp/post-path\nexit $rv\n}\n" +
+			"trap -- \"__cleanup\" EXIT\n"
+		require.EqualValues(t, expected, buf.String())
+	})
 }
 
 func TestSetOnShell_SkipShellHistory(t *testing.T) {