From b46a75d8024c33d393e74f00860bfe9c25e209ed Mon Sep 17 00:00:00 2001 From: Adam Babik Date: Sun, 13 Oct 2024 15:19:51 +0200 Subject: [PATCH] Update unit tests for env collectors --- internal/cmd/beta/session_cmd.go | 2 +- internal/command/env_collector_fifo_unix.go | 15 ++++--- .../command/env_collector_fifo_unix_test.go | 29 ++++++++++--- internal/command/env_collector_file.go | 15 ++++--- internal/command/env_collector_file_test.go | 27 +++++++++--- internal/command/env_shell.go | 2 +- internal/command/env_shell_test.go | 41 ++++++++++++++----- 7 files changed, 95 insertions(+), 36 deletions(-) diff --git a/internal/cmd/beta/session_cmd.go b/internal/cmd/beta/session_cmd.go index b92287b57..fe6e2ec69 100644 --- 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 index cb112e024..73aa60318 100644 --- 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 index 3cf04ae3a..93383b98a 100644 --- 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 index a0124869b..5fdab47a2 100644 --- 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 index 31199072d..739783ab9 100644 --- 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 index 7898a9626..9560713d9 100644 --- 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 index e7cd0bab1..3fa06b5be 100644 --- 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) {