refactor: mv other envs to processEnvVars

replicatedhq · Oct 17, 2023 · 48ee607 · 48ee607
1 parent 632b5c0
commit 48ee607
Showing 1 changed file with 55 additions and 50 deletions.
diff --git a/pkg/collect/host_run.go b/pkg/collect/host_run.go
@@ -43,8 +43,6 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][]
 	}
 	var (
 		bundleOutputRelativePath string
-		cmdOutputTempDir         string
-		cmdInputTempDir          string
 		err                      error
 	)
 
@@ -55,47 +53,14 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][]
 		ExitCode: "0",
 	}
 
-	err = c.processEnvVars(cmd, runInfo)
+	cmdOutputTempDir, err := c.processEnvVars(cmd)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to parse env variable")
 	}
 
 	collectorRelativePath := filepath.Join("host-collectors/run-host", collectorName)
 
-	// if we choose to save result for the command run
-	if runHostCollector.OutputDir != "" {
-		cmdOutputTempDir, err = os.MkdirTemp("", runHostCollector.OutputDir)
-		defer os.RemoveAll(cmdOutputTempDir)
-		if err != nil {
-			return nil, errors.New(fmt.Sprintf("failed to created temp dir for: %s", runHostCollector.OutputDir))
-		}
-		cmd.Env = append(cmd.Env,
-			fmt.Sprintf("TS_WORKSPACE_DIR=%s", cmdOutputTempDir),
-		)
-	}
-
-	output := NewResult()
-
-	if runHostCollector.Input != nil {
-		cmdInputTempDir, err = os.MkdirTemp("", "input")
-		defer os.RemoveAll(cmdInputTempDir)
-		if err != nil {
-			return nil, errors.New("failed to created temp dir for host run input")
-		}
-		for inFilename, inFileContent := range runHostCollector.Input {
-			if strings.Contains(inFileContent, "/") {
-				return nil, errors.New("Input filename contains '/'")
-			}
-			cmdInputFilePath := filepath.Join(cmdInputTempDir, inFilename)
-			err = os.WriteFile(cmdInputFilePath, []byte(inFileContent), 0644)
-			if err != nil {
-				return nil, errors.Wrap(err, fmt.Sprintf("failed to write input file: %s to temp directory", inFilename))
-			}
-		}
-		cmd.Env = append(cmd.Env,
-			fmt.Sprintf("TS_INPUT_DIR=%s", cmdInputTempDir),
-		)
-	}
+	runInfo.Env = cmd.Env
 
 	var stdout, stderr bytes.Buffer
 	cmd.Stdout = &stdout
@@ -111,6 +76,7 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][]
 		}
 	}
 
+	output := NewResult()
 	resultInfo := filepath.Join("host-collectors/run-host", collectorName+"-info.json")
 	result := filepath.Join("host-collectors/run-host", collectorName+".txt")
 
@@ -131,7 +97,13 @@ func (c *CollectHostRun) Collect(progressChan chan<- interface{}) (map[string][]
 	return output, nil
 }
 
-func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd, runInfo *HostRunInfo) error {
+func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd) (string, error) {
+	var (
+		cmdOutputTempDir string
+		err              error
+	)
+	// clears the parent env vars
+	cmd.Env = []string{}
 	guaranteedEnvs := []string{"PATH", "KUBECONFIG"}
 	for _, key := range guaranteedEnvs {
 		guaranteedEnvVal, found := os.LookupEnv(key)
@@ -142,16 +114,18 @@ func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd, runInfo *HostRunInfo) err
 	}
 
 	runHostCollector := c.hostCollector
-	if runHostCollector.InheritParentEnvs == true {
-		cmd.Env = append(cmd.Env, os.Environ()...)
-	} else if runHostCollector.InheritEnvs != nil {
-		for _, key := range runHostCollector.InheritEnvs {
-			envVal, found := os.LookupEnv(key)
-			if !found {
-				return errors.New(fmt.Sprintf("inherit env variable is not found: %s", key))
+	if runHostCollector.InheritParentEnvs {
+		if runHostCollector.InheritEnvs != nil {
+			for _, key := range runHostCollector.InheritEnvs {
+				envVal, found := os.LookupEnv(key)
+				if !found {
+					return "", errors.New(fmt.Sprintf("inherit env variable is not found: %s", key))
+				}
+				cmd.Env = append(cmd.Env,
+					fmt.Sprintf("%s=%s", key, envVal))
 			}
-			cmd.Env = append(cmd.Env,
-				fmt.Sprintf("%s=%s", key, envVal))
+		} else {
+			cmd.Env = append(cmd.Env, os.Environ()...)
 		}
 	}
 
@@ -162,12 +136,43 @@ func (c *CollectHostRun) processEnvVars(cmd *exec.Cmd, runInfo *HostRunInfo) err
 				cmd.Env = append(cmd.Env,
 					fmt.Sprintf("%s=%s", runHostCollector.Env[i], os.Getenv(runHostCollector.Env[i])))
 			} else {
-				return errors.New(fmt.Sprintf("env variable entry is missing '=' : %s", runHostCollector.Env[i]))
+				return "", errors.New(fmt.Sprintf("env variable entry is missing '=' : %s", runHostCollector.Env[i]))
 			}
 		}
 	}
 
-	runInfo.Env = os.Environ()
+	// if we choose to save result for the command run
+	if runHostCollector.OutputDir != "" {
+		cmdOutputTempDir, err = os.MkdirTemp("", runHostCollector.OutputDir)
+		defer os.RemoveAll(cmdOutputTempDir)
+		if err != nil {
+			return "", errors.New(fmt.Sprintf("failed to created temp dir for: %s", runHostCollector.OutputDir))
+		}
+		cmd.Env = append(cmd.Env,
+			fmt.Sprintf("TS_WORKSPACE_DIR=%s", cmdOutputTempDir),
+		)
+	}
+
+	if runHostCollector.Input != nil {
+		cmdInputTempDir, err := os.MkdirTemp("", "input")
+		defer os.RemoveAll(cmdInputTempDir)
+		if err != nil {
+			return "", errors.New("failed to created temp dir for host run input")
+		}
+		for inFilename, inFileContent := range runHostCollector.Input {
+			if strings.Contains(inFileContent, "/") {
+				return "", errors.New("Input filename contains '/'")
+			}
+			cmdInputFilePath := filepath.Join(cmdInputTempDir, inFilename)
+			err = os.WriteFile(cmdInputFilePath, []byte(inFileContent), 0644)
+			if err != nil {
+				return "", errors.Wrap(err, fmt.Sprintf("failed to write input file: %s to temp directory", inFilename))
+			}
+		}
+		cmd.Env = append(cmd.Env,
+			fmt.Sprintf("TS_INPUT_DIR=%s", cmdInputTempDir),
+		)
+	}
 
-	return nil
+	return cmdOutputTempDir, nil
 }