From 5e67c56512ad235b1a79ff16dbbedb95f095f811 Mon Sep 17 00:00:00 2001
From: Shantanu <joshishantanu4@gmail.com>
Date: Fri, 3 May 2024 15:03:36 -0700
Subject: [PATCH] shell/zsh: Setup env vars for runner shell in one func.

The runbookRunMetadata func is used to setup both zsh and bash shells
---
 shell/zsh.go | 43 ++++++++++++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 15 deletions(-)

diff --git a/shell/zsh.go b/shell/zsh.go
index eba88c7..a106fa9 100644
--- a/shell/zsh.go
+++ b/shell/zsh.go
@@ -246,22 +246,9 @@ func (z *zsh) SpawnRunbookRunner(ctx context.Context, runbook *client.Runbook) (
 		return nil, err
 	}
 
-	// Inherit the next step from the environment if we are in a subshell
-	nextStep := os.Getenv("SAVVY_NEXT_STEP")
-	if nextStep == "" {
-		nextStep = "1"
-	}
-
-	var nextStepIdx int
-	nextStepIdx, err = strconv.Atoi(nextStep)
-	if err != nil {
-		nextStepIdx = 1
-	}
-
-	runbook_alias := computeRunbookAlias(runbook)
-
 	cmd := exec.CommandContext(ctx, z.shellCmd)
-	cmd.Env = append(os.Environ(), "ZDOTDIR="+tmp, "SAVVY_CONTEXT=run", fmt.Sprintf("SAVVY_RUNBOOK_COMMANDS=%s", strings.Join(runbook.Commands(), RunbookCommandDelimiter)), fmt.Sprintf("SAVVY_NEXT_STEP=%d", nextStepIdx), fmt.Sprintf("SAVVY_RUNBOOK_ALIAS=%s", runbook_alias))
+	cmd.Env = append(os.Environ(), "ZDOTDIR="+tmp)
+	cmd.Env = append(cmd.Env, runbookRunMetadata(runbook)...)
 	cmd.WaitDelay = 500 * time.Millisecond
 	return cmd, nil
 }
@@ -272,3 +259,29 @@ func computeRunbookAlias(runbook *client.Runbook) string {
 	alias, _ = strings.CutPrefix(alias, "how-to-")
 	return alias
 }
+
+func nextRunbookStepToRun() int {
+	// Inherit the next step from the environment if we are in a subshell
+	nextStep := os.Getenv("SAVVY_NEXT_STEP")
+	if nextStep == "" {
+		nextStep = "1"
+	}
+
+	idx, err := strconv.Atoi(nextStep)
+	if err != nil {
+		return 1
+	}
+	return idx
+}
+
+func runbookRunMetadata(runbook *client.Runbook) []string {
+	runbookCommands := strings.Join(runbook.Commands(), RunbookCommandDelimiter)
+	runbookAlias := computeRunbookAlias(runbook)
+
+	return []string{
+		"SAVVY_CONTEXT=run",
+		fmt.Sprintf("SAVVY_RUNBOOK_COMMANDS=%s", runbookCommands),
+		fmt.Sprintf("SAVVY_NEXT_STEP=%d", nextRunbookStepToRun()),
+		fmt.Sprintf("SAVVY_RUNBOOK_ALIAS=%s", runbookAlias),
+	}
+}