Create mutable variable store (#514)

Summary: Pull Request resolved: #514 Create a variable store that we can mutate between test steps. Move WorkDir into this variable store so that we can mutate the work dir for a future cd step. In the future if we want to carry any other information between steps in a mutable fashion, all we need to do is add a new variable to this struct. Reviewed By: inesusvet Differential Revision: D63457210
facebookincubator · Oct 1, 2024 · 904422a · 904422a
1 parent 62814c1
commit 904422a
Show file tree

Hide file tree

Showing 10 changed files with 35 additions and 19 deletions.
diff --git a/pkg/blocks/context.go b/pkg/blocks/context.go
@@ -41,10 +41,15 @@ type TTPExecutionConfig struct {
 	Stderr              io.Writer
 }
 
+// TTPExecutionVars - mutable store to carry variables between steps
+type TTPExecutionVars struct {
+	WorkDir string
+}
+
 // TTPExecutionContext - holds config and context for the currently executing TTP
 type TTPExecutionContext struct {
 	Cfg               TTPExecutionConfig
-	WorkDir           string
+	Vars              *TTPExecutionVars
 	StepResults       *StepResultsRecord
 	actionResultsChan chan *ActResult
 	errorsChan        chan error
@@ -54,6 +59,7 @@ type TTPExecutionContext struct {
 // NewTTPExecutionContext creates a new TTPExecutionContext with empty config and created channels
 func NewTTPExecutionContext() TTPExecutionContext {
 	return TTPExecutionContext{
+		Vars:              &TTPExecutionVars{WorkDir: "/"},
 		StepResults:       NewStepResultsRecord(),
 		actionResultsChan: make(chan *ActResult, 1),
 		errorsChan:        make(chan error, 1),

diff --git a/pkg/blocks/editstep.go b/pkg/blocks/editstep.go
@@ -71,7 +71,7 @@ func (s *EditStep) Validate(execCtx TTPExecutionContext) error {
 	fileSystem := s.FileSystem
 
 	if fileSystem == nil {
-		_, err := FetchAbs(targetPath, execCtx.WorkDir)
+		_, err := FetchAbs(targetPath, execCtx.Vars.WorkDir)
 		if err != nil {
 			return err
 		}
@@ -130,12 +130,12 @@ func (s *EditStep) Execute(execCtx TTPExecutionContext) (*ActResult, error) {
 	if fileSystem == nil {
 		fileSystem = afero.NewOsFs()
 		var err error
-		targetPath, err = FetchAbs(targetPath, execCtx.WorkDir)
+		targetPath, err = FetchAbs(targetPath, execCtx.Vars.WorkDir)
 		if err != nil {
 			return nil, err
 		}
 		if backupPath != "" {
-			backupPath, err = FetchAbs(backupPath, execCtx.WorkDir)
+			backupPath, err = FetchAbs(backupPath, execCtx.Vars.WorkDir)
 			if err != nil {
 				return nil, err
 			}

diff --git a/pkg/blocks/editstep_test.go b/pkg/blocks/editstep_test.go
@@ -261,7 +261,7 @@ edits:
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			var editStep EditStep
-			var execCtx TTPExecutionContext
+			execCtx := NewTTPExecutionContext()
 
 			// parse the step
 			err := yaml.Unmarshal([]byte(tc.content), &editStep)

diff --git a/pkg/blocks/executor.go b/pkg/blocks/executor.go
@@ -108,7 +108,7 @@ func (e *ScriptExecutor) Execute(ctx context.Context, execCtx TTPExecutionContex
 
 	cmd := e.buildCommand(ctx)
 	cmd.Env = expandedEnvAsList
-	cmd.Dir = execCtx.WorkDir
+	cmd.Dir = execCtx.Vars.WorkDir
 	cmd.Stdin = strings.NewReader(body)
 
 	return streamAndCapture(*cmd, execCtx.Cfg.Stdout, execCtx.Cfg.Stderr)
@@ -138,7 +138,7 @@ func (e *FileExecutor) Execute(ctx context.Context, execCtx TTPExecutionContext)
 	}
 
 	cmd.Env = expandedEnvAsList
-	cmd.Dir = execCtx.WorkDir
+	cmd.Dir = execCtx.Vars.WorkDir
 	return streamAndCapture(*cmd, execCtx.Cfg.Stdout, execCtx.Cfg.Stderr)
 }
 

diff --git a/pkg/blocks/expectstep.go b/pkg/blocks/expectstep.go
@@ -258,7 +258,7 @@ func (s *ExpectStep) prepareCommand(ctx context.Context, execCtx TTPExecutionCon
 	/* #nosec G204 */
 	cmd := exec.CommandContext(ctx, s.Executor, "-c", inline)
 	cmd.Env = envAsList
-	cmd.Dir = execCtx.WorkDir
+	cmd.Dir = execCtx.Vars.WorkDir
 
 	return cmd
 }

diff --git a/pkg/blocks/expectstep_test.go b/pkg/blocks/expectstep_test.go
@@ -57,7 +57,9 @@ func createTestScript(t *testing.T, scriptContent string) (string, string) {
 
 func NewTestTTPExecutionContext(workDir string) TTPExecutionContext {
 	return TTPExecutionContext{
-		WorkDir: workDir,
+		Vars: &TTPExecutionVars{
+			WorkDir: workDir,
+		},
 	}
 }
 
@@ -668,7 +670,11 @@ func TestExpectSSH(t *testing.T) {
 			mockStep := new(MockExpectStep)
 			mockStep.On("Execute", mock.Anything).Return(tc.mockRet, tc.mockErr)
 
-			execCtx := TTPExecutionContext{WorkDir: "."}
+			execCtx := TTPExecutionContext{
+				Vars: &TTPExecutionVars{
+					WorkDir: ".",
+				},
+			}
 			fmt.Println("Executing command:", tc.step.Expect.Inline)
 			_, err := mockStep.Execute(execCtx)
 			if (err != nil) != tc.wantErr {
@@ -744,7 +750,9 @@ func TestCleanup(t *testing.T) {
 			}
 
 			execCtx := TTPExecutionContext{
-				WorkDir: ".",
+				Vars: &TTPExecutionVars{
+					WorkDir: ".",
+				},
 			}
 
 			_, err := s.Cleanup(execCtx)

diff --git a/pkg/blocks/fetchuri.go b/pkg/blocks/fetchuri.go
@@ -91,7 +91,7 @@ func (f *FetchURIStep) Validate(execCtx TTPExecutionContext) error {
 	}
 
 	// Retrieve the absolute path to the file.
-	absLocal, err := FetchAbs(f.Location, execCtx.WorkDir)
+	absLocal, err := FetchAbs(f.Location, execCtx.Vars.WorkDir)
 	if err != nil {
 		logging.L().Error(zap.Error(err))
 		return err
@@ -135,7 +135,7 @@ func (f *FetchURIStep) fetchURI(execCtx TTPExecutionContext) error {
 	if appFs == nil {
 		var err error
 		appFs = afero.NewOsFs()
-		absLocal, err = FetchAbs(f.Location, execCtx.WorkDir)
+		absLocal, err = FetchAbs(f.Location, execCtx.Vars.WorkDir)
 		if err != nil {
 			return err
 		}

diff --git a/pkg/blocks/fetchuri_test.go b/pkg/blocks/fetchuri_test.go
@@ -181,7 +181,7 @@ steps:
 			err := yaml.Unmarshal([]byte(tc.content), &ttps)
 			assert.NoError(t, err)
 
-			var execCtx TTPExecutionContext
+			execCtx := NewTTPExecutionContext()
 			err = ttps.Validate(execCtx)
 			if tc.wantError {
 				assert.Error(t, err)
@@ -202,7 +202,7 @@ location: /tmp/test.html
 overwrite: true
 `
 	var s FetchURIStep
-	var execCtx TTPExecutionContext
+	execCtx := NewTTPExecutionContext()
 	err := yaml.Unmarshal([]byte(content), &s)
 	require.NoError(t, err)
 

diff --git a/pkg/blocks/filestep.go b/pkg/blocks/filestep.go
@@ -74,14 +74,14 @@ func (f *FileStep) Validate(execCtx TTPExecutionContext) error {
 	}
 
 	// If FilePath is set, ensure that the file exists.
-	fullpath, err := FindFilePath(f.FilePath, execCtx.WorkDir, nil)
+	fullpath, err := FindFilePath(f.FilePath, execCtx.Vars.WorkDir, nil)
 	if err != nil {
 		logging.L().Error(zap.Error(err))
 		return err
 	}
 
 	// Retrieve the absolute path to the file.
-	f.FilePath, err = FetchAbs(fullpath, execCtx.WorkDir)
+	f.FilePath, err = FetchAbs(fullpath, execCtx.Vars.WorkDir)
 	if err != nil {
 		logging.L().Error(zap.Error(err))
 		return err

diff --git a/pkg/blocks/loader.go b/pkg/blocks/loader.go
@@ -151,8 +151,10 @@ func LoadTTP(ttpFilePath string, fsys afero.Fs, execCfg *TTPExecutionConfig, arg
 	}
 
 	execCtx := TTPExecutionContext{
-		Cfg:               *execCfg,
-		WorkDir:           ttp.WorkDir,
+		Cfg: *execCfg,
+		Vars: &TTPExecutionVars{
+			WorkDir: ttp.WorkDir,
+		},
 		StepResults:       NewStepResultsRecord(),
 		actionResultsChan: make(chan *ActResult, 1),
 		errorsChan:        make(chan error, 1),