From 31b9365d9d20a766c18dc00f5d4c33d5abfd4c70 Mon Sep 17 00:00:00 2001
From: Abhinav Gupta <abgupta@rippling.com>
Date: Mon, 19 Feb 2024 10:37:19 -0800
Subject: [PATCH] testscript: Add Chdir method to change directory

It is not currently possible for a custom testscript command
to change the working directory of that script run.
`TestScript.Exec` runs the command in a subprocess,
so one cannot do `ts.Exec("cd", dir)`.

This change adds a `Chdir` method to `TestScript`
that allows changing the working directory of the script.
The implementation is the same as the "cd" command,
which now relies on `Chdir`.

The availability of this function matches similar functionality in the
[`State.Chdir` method of rsc.io/script][1].
(I ported some tests from rsc.io/script to testscript.)

  [1]: https://pkg.go.dev/rsc.io/script#State.Chdir
---
 testscript/cmd.go                 |  9 +--------
 testscript/testdata/custom_cd.txt |  7 +++++++
 testscript/testscript.go          | 18 ++++++++++++++++++
 testscript/testscript_test.go     | 21 +++++++++++++++++++++
 4 files changed, 47 insertions(+), 8 deletions(-)
 create mode 100644 testscript/testdata/custom_cd.txt

diff --git a/testscript/cmd.go b/testscript/cmd.go
index 5e450766..446245a1 100644
--- a/testscript/cmd.go
+++ b/testscript/cmd.go
@@ -61,18 +61,11 @@ func (ts *TestScript) cmdCd(neg bool, args []string) {
 	}
 
 	dir := args[0]
-	if !filepath.IsAbs(dir) {
-		dir = filepath.Join(ts.cd, dir)
-	}
-	info, err := os.Stat(dir)
+	err := ts.Chdir(dir)
 	if os.IsNotExist(err) {
 		ts.Fatalf("directory %s does not exist", dir)
 	}
 	ts.Check(err)
-	if !info.IsDir() {
-		ts.Fatalf("%s is not a directory", dir)
-	}
-	ts.cd = dir
 	ts.Logf("%s\n", ts.cd)
 }
 
diff --git a/testscript/testdata/custom_cd.txt b/testscript/testdata/custom_cd.txt
new file mode 100644
index 00000000..05d9689e
--- /dev/null
+++ b/testscript/testdata/custom_cd.txt
@@ -0,0 +1,7 @@
+# Verify that a custom command can chdir.
+
+mkChdir foo
+exists $WORK/foo
+
+# Current directory is not $WORK.
+! exists foo
diff --git a/testscript/testscript.go b/testscript/testscript.go
index 089e1038..a6b3e5c8 100644
--- a/testscript/testscript.go
+++ b/testscript/testscript.go
@@ -1043,6 +1043,24 @@ func (ts *TestScript) BackgroundCmds() []*exec.Cmd {
 	return cmds
 }
 
+// Chdir changes the current directory of the script.
+// The path may be relative to the current directory.
+func (ts *TestScript) Chdir(dir string) error {
+	if !filepath.IsAbs(dir) {
+		dir = filepath.Join(ts.cd, dir)
+	}
+	info, err := os.Stat(dir)
+	if err != nil {
+		return err
+	}
+	if !info.IsDir() {
+		return fmt.Errorf("%s is not a directory", dir)
+	}
+
+	ts.cd = dir
+	return nil
+}
+
 // waitOrStop waits for the already-started command cmd by calling its Wait method.
 //
 // If cmd does not return before ctx is done, waitOrStop sends it an interrupt
diff --git a/testscript/testscript_test.go b/testscript/testscript_test.go
index 4a62dbb7..f01b468f 100644
--- a/testscript/testscript_test.go
+++ b/testscript/testscript_test.go
@@ -263,6 +263,27 @@ func TestScripts(t *testing.T) {
 				}
 			},
 			"echoandexit": echoandexit,
+			"mkChdir": func(ts *TestScript, neg bool, args []string) {
+				if neg {
+					ts.Fatalf("unsupported: ! mkChdir")
+				}
+				if len(args) != 1 {
+					ts.Fatalf("usage: mkChdir dir")
+				}
+
+				dir := args[0]
+				if !filepath.IsAbs(dir) {
+					dir = ts.MkAbs(dir)
+				}
+
+				if err := os.MkdirAll(dir, 0o777); err != nil {
+					ts.Fatalf("cannot create dir: %v", err)
+				}
+
+				if err := ts.Chdir(dir); err != nil {
+					ts.Fatalf("cannot chdir: %v", err)
+				}
+			},
 		},
 		Setup: func(env *Env) error {
 			infos, err := os.ReadDir(env.WorkDir)