redirect_tty.go

Documentation: github.com/dsoprea/go-utility/v2/testing

     1  package ritesting
     2  
     3  import (
     4  	"fmt"
     5  	"os"
     6  
     7  	"io/ioutil"
     8  
     9  	"github.com/dsoprea/go-logging"
    10  )
    11  
    12  type redirectedTty struct {
    13  	originalStdin  *os.File
    14  	newStdinWriter *os.File
    15  
    16  	originalStdout  *os.File
    17  	newStdoutReader *os.File
    18  
    19  	originalStderr  *os.File
    20  	newStderrReader *os.File
    21  }
    22  
    23  var (
    24  	rtty *redirectedTty
    25  )
    26  
    27  // RedirectTty redirects STDIN, STDOUT, and STDERR to alternative resources.
    28  // This allows us to create unit-tests for executables by directly calling
    29  // their main() entrypoints.
    30  func RedirectTty() {
    31  	if rtty != nil {
    32  		log.Panicf("TTY is already redirected")
    33  	}
    34  
    35  	rtty = &redirectedTty{
    36  		originalStdin:  os.Stdin,
    37  		originalStdout: os.Stdout,
    38  		originalStderr: os.Stderr,
    39  	}
    40  
    41  	var err error
    42  
    43  	// The caller can write to the new writer and it will be available on STDIN.
    44  
    45  	os.Stdin, rtty.newStdinWriter, err = os.Pipe()
    46  	log.PanicIf(err)
    47  
    48  	// Any output to STDOUT or STDERR will be available to the caller on the
    49  	// new readers.
    50  
    51  	rtty.newStdoutReader, os.Stdout, err = os.Pipe()
    52  	log.PanicIf(err)
    53  
    54  	rtty.newStderrReader, os.Stderr, err = os.Pipe()
    55  	log.PanicIf(err)
    56  }
    57  
    58  // RestoreTty restores original TTY resources.
    59  func RestoreTty() {
    60  	if rtty == nil {
    61  		return
    62  	}
    63  
    64  	os.Stdin = rtty.originalStdin
    65  	os.Stdout = rtty.originalStdout
    66  	os.Stderr = rtty.originalStderr
    67  
    68  	rtty = nil
    69  }
    70  
    71  // RestoreAndDumpTty restores original TTY resources but not before grabbing
    72  // their screen output and then printing it before returning.
    73  func RestoreAndDumpTty() {
    74  	if rtty == nil {
    75  		return
    76  	}
    77  
    78  	// TODO(dustin): !! Finish. Close os.Stdout and os.Stderr, read each, and print each between anchors.
    79  
    80  	os.Stdout.Close()
    81  
    82  	stdoutOutput, err := ioutil.ReadAll(rtty.newStdoutReader)
    83  	log.PanicIf(err)
    84  
    85  	os.Stderr.Close()
    86  
    87  	stderrOutput, err := ioutil.ReadAll(rtty.newStderrReader)
    88  	log.PanicIf(err)
    89  
    90  	RestoreTty()
    91  
    92  	fmt.Printf(">>>>>>>>>>>>>\n")
    93  	fmt.Printf("STDOUT OUTPUT\n")
    94  	fmt.Printf(">>>>>>>>>>>>>\n")
    95  	fmt.Println(string(stdoutOutput))
    96  	fmt.Printf("<<<<<<<<<<<<<\n")
    97  	fmt.Printf("\n")
    98  
    99  	fmt.Printf(">>>>>>>>>>>>>\n")
   100  	fmt.Printf("STDERR OUTPUT\n")
   101  	fmt.Printf(">>>>>>>>>>>>>\n")
   102  	fmt.Println(string(stderrOutput))
   103  	fmt.Printf("<<<<<<<<<<<<<\n")
   104  	fmt.Printf("\n")
   105  }
   106  
   107  // StdinWriter returns a writer that can be used to feed STDIN data (if
   108  // redirected).
   109  func StdinWriter() *os.File {
   110  	if rtty == nil {
   111  		log.Panicf("TTY not redirected; STDIN writer not available")
   112  	}
   113  
   114  	return rtty.newStdinWriter
   115  }
   116  
   117  // StdoutReader returns a reader that can be used to read STDOUT output (if
   118  // redirected).
   119  func StdoutReader() *os.File {
   120  	if rtty == nil {
   121  		log.Panicf("TTY not redirected; STDOUT reader not available")
   122  	}
   123  
   124  	return rtty.newStdoutReader
   125  }
   126  
   127  // StderrReader returns a reader that can be used to read STDERR output (if
   128  // redirected).
   129  func StderrReader() *os.File {
   130  	if rtty == nil {
   131  		log.Panicf("TTY not redirected; STDERR reader not available")
   132  	}
   133  
   134  	return rtty.newStderrReader
   135  }
   136  
   137  // IsTtyRedirected returns whether the TTY is currently redirected.
   138  func IsTtyRedirected() bool {
   139  	return rtty != nil
   140  }
   141
View as plain text