cmd/testscript: support -update flag

This is exactly equivalent to setting testscript.Params.UpdateScripts

cmd/testscript/main.go

@@ -62,6 +62,7 @@ func mainerr() (retErr error) {
 		mainUsage(os.Stderr)
 	}
 	var envVars envVarsFlag
+	fUpdate := fs.Bool("update", false, "update archive file if a cmp fails")
 	fWork := fs.Bool("work", false, "print temporary work directory and do not remove when done")
 	fVerbose := fs.Bool("v", false, "run tests verbosely")
 	fs.Var(&envVars, "e", "pass through environment variable to script (can appear multiple times)")
@@ -103,7 +104,7 @@ func mainerr() (retErr error) {
 		if err := os.Mkdir(runDir, 0777); err != nil {
 			return fmt.Errorf("failed to create a run directory within %v for %v: %v", td, renderFilename(filename), err)
 		}
-		if err := run(runDir, filename, *fVerbose, envVars.vals); err != nil {
+		if err := run(runDir, filename, *fUpdate, *fVerbose, envVars.vals); err != nil {
 			return err
 		}
 	}
@@ -162,7 +163,11 @@ func renderFilename(filename string) string {
 	return filename
 }
 
-func run(runDir, filename string, verbose bool, envVars []string) error {
+// run runs the testscript archive in filename within the temporary runDir.
+// verbose causes the output to be verbose (akin to go test -v) and update
+// sets the UpdateScripts parameter passed to testscript.Run such that any
+// updates to the archive get written back to filename
+func run(runDir, filename string, update bool, verbose bool, envVars []string) error {
 	var ar *txtar.Archive
 	var err error
 
@@ -204,12 +209,15 @@ func run(runDir, filename string, verbose bool, envVars []string) error {
 		return fmt.Errorf("failed to write .gomodproxy files: %v", err)
 	}
 
-	if err := ioutil.WriteFile(filepath.Join(runDir, "script.txt"), txtar.Format(&script), 0666); err != nil {
-		return fmt.Errorf("failed to write script for %v: %v", filename, err)
+	scriptFile := filepath.Join(runDir, "script.txt")
+
+	if err := ioutil.WriteFile(scriptFile, txtar.Format(&script), 0666); err != nil {
+		return fmt.Errorf("failed to write script for %v: %v", renderFilename(filename), err)
 	}
 
 	p := testscript.Params{
-		Dir: runDir,
+		Dir:           runDir,
+		UpdateScripts: update,
 	}
 
 	if _, err := exec.LookPath("go"); err == nil {
@@ -282,5 +290,16 @@ func run(runDir, filename string, verbose bool, envVars []string) error {
 		return fmt.Errorf("error running %v in %v\n", renderFilename(filename), runDir)
 	}
 
+	if update {
+		// Write the testscript archive back to the source file
+		source, err := ioutil.ReadFile(scriptFile)
+		if err != nil {
+			return fmt.Errorf("failed to read from script file %v for -update: %v", scriptFile, err)
+		}
+		if err := ioutil.WriteFile(filename, source, 0666); err != nil {
+			return fmt.Errorf("failed to write script back to %v for -update: %v", renderFilename(filename), err)
+		}
+	}
+
 	return nil
 }

cmd/testscript/testdata/update.txt

@@ -0,0 +1,19 @@
+# should support the -update flag
+
+unquote in.txt res.txt
+
+testscript -update in.txt
+cmp in.txt res.txt
+
+-- in.txt --
+>exec printf 'hello\n'
+>cmp stdout stdout.txt
+>
+>-- stdout.txt --
+>goodbyte
+-- res.txt --
+>exec printf 'hello\n'
+>cmp stdout stdout.txt
+>
+>-- stdout.txt --
+>hello