github.com/ActiveState/cli@v0.0.0-20240508170324-6801f60cd051/internal/testhelpers/osutil/std_wrap.go (about)

     1  package osutil
     2  
     3  import (
     4  	"io"
     5  	"log"
     6  	"os"
     7  	"time"
     8  
     9  	"github.com/fatih/color"
    10  )
    11  
    12  func replaceStderr(newErr *os.File) *os.File {
    13  	oldErr := os.Stderr
    14  	os.Stderr = newErr
    15  	color.Output = newErr
    16  	return oldErr
    17  }
    18  
    19  func replaceStdout(newOut *os.File) *os.File {
    20  	oldOut := os.Stdout
    21  	os.Stdout = newOut
    22  	color.Output = newOut
    23  	return oldOut
    24  }
    25  
    26  func replaceStdin(newIn *os.File) *os.File {
    27  	oldIn := os.Stdin
    28  	os.Stdin = newIn
    29  	return oldIn
    30  }
    31  
    32  // captureWrites will execute a provided function and return any bytes written to the provided
    33  // writer from the provided reader (assuming they are generated from something like an os.Pipe).
    34  func captureWrites(fnToExec func(), reader, writer *os.File) (string, error) {
    35  	fnToExec() // execute the provided function
    36  
    37  	if err := writer.Close(); err != nil {
    38  		return "", err
    39  	}
    40  
    41  	writeBytes, err := io.ReadAll(reader)
    42  	if err != nil {
    43  		err = reader.Close()
    44  	}
    45  	return string(writeBytes), err
    46  }
    47  
    48  // CaptureStderr will execute a provided function and capture anything written to stderr.
    49  // It will then return that output as a string along with any errors captured in the process.
    50  func CaptureStderr(fnToExec func()) (string, error) {
    51  	errReader, errWriter, err := os.Pipe()
    52  	if err != nil {
    53  		return "", err
    54  	}
    55  	defer replaceStderr(replaceStderr(errWriter))
    56  	return captureWrites(fnToExec, errReader, errWriter)
    57  }
    58  
    59  // CaptureStdout will execute a provided function and capture anything written to stdout.
    60  // It will then return that output as a string along with any errors captured in the process.
    61  func CaptureStdout(fnToExec func()) (string, error) {
    62  	outReader, outWriter, err := os.Pipe()
    63  	if err != nil {
    64  		return "", err
    65  	}
    66  	defer replaceStdout(replaceStdout(outWriter))
    67  	return captureWrites(fnToExec, outReader, outWriter)
    68  }
    69  
    70  // WrapStdin will fill stdin with the lines provided as a variadic list of strings before
    71  // executing the provided function. Each line will be appended with a newline (\n) only.
    72  func WrapStdin(fnToExec func(), inputLines ...interface{}) {
    73  	WrapStdinWithDelay(0, fnToExec, inputLines...)
    74  }
    75  
    76  // WrapStdinWithDelay will fill stdin with the lines provided, but with a given delay before
    77  // each write. This is useful if there is a reader that reads all of stdin between prompts,
    78  // for instance.
    79  func WrapStdinWithDelay(delay time.Duration, fnToExec func(), inputLines ...interface{}) {
    80  	tmpIn, inWriter, err := os.Pipe()
    81  	if err != nil {
    82  		panic(err)
    83  	}
    84  	defer tmpIn.Close()
    85  	defer replaceStdin(replaceStdin(tmpIn))
    86  
    87  	if delay > 0 {
    88  		// need to run this asynchornously so that the fnToExec can be processed
    89  		go func() {
    90  			time.Sleep(500 * time.Millisecond) // Give fnToExec some time to start
    91  			writeLinesAndClosePipe(inWriter, inputLines, func() { time.Sleep(delay) })
    92  		}()
    93  	} else {
    94  		writeLinesAndClosePipe(inWriter, inputLines, nil)
    95  	}
    96  
    97  	fnToExec() // execute the provided function
    98  }
    99  
   100  func writeLinesAndClosePipe(writer *os.File, lines []interface{}, callbackFn func()) {
   101  	defer writer.Close()
   102  	for _, line := range lines {
   103  		if callbackFn != nil {
   104  			callbackFn()
   105  		}
   106  		if lineStr, ok := line.(string); ok {
   107  			_, err := writer.WriteString(lineStr + "\n")
   108  			if err != nil {
   109  				log.Panicf("Error writing to stdin: %v", err)
   110  			}
   111  		} else if lineRune, ok := line.(rune); ok {
   112  			_, err := writer.WriteString(string(lineRune))
   113  			if err != nil {
   114  				log.Panicf("Error writing to stdin: %v", err)
   115  			}
   116  		} else {
   117  			log.Panicf("Unsupported line: %v", line)
   118  		}
   119  	}
   120  }