github.com/elopio/cli@v6.21.2-0.20160902224010-ea909d1fdb2f+incompatible/testhelpers/io/io.go (about)

     1  package io
     2  
     3  import (
     4  	"bytes"
     5  	"github.com/fatih/color"
     6  	"github.com/mattn/go-colorable"
     7  	"io"
     8  	"os"
     9  	"runtime"
    10  	"strings"
    11  )
    12  
    13  func SimulateStdin(input string, block func(r io.Reader)) {
    14  	reader, writer := io.Pipe()
    15  
    16  	go func() {
    17  		writer.Write([]byte(input))
    18  		defer writer.Close()
    19  	}()
    20  
    21  	block(reader)
    22  }
    23  
    24  func CaptureOutput(block func()) []string {
    25  	oldSTDOUT := os.Stdout
    26  	r, w, err := os.Pipe()
    27  	if err != nil {
    28  		panic(err)
    29  	}
    30  
    31  	os.Stdout = w
    32  	defer func() {
    33  		os.Stdout = oldSTDOUT
    34  	}()
    35  
    36  	//////
    37  	// We use fmt.Fprintf() to write to the "github.com/fatih/color".Output file
    38  	// to get colors on Windows machines.
    39  	// That variable gets initialized with a reference to os.Stdout when that library is imported.
    40  	// That means that when we muck with os.Stdout above, it doesn't get reflected in
    41  	// the printing code for windows.
    42  	// Instead, we can just redeclare that color.Output variable with a colorable version of our
    43  	// redirect pipe.
    44  	if runtime.GOOS == "windows" {
    45  		color.Output = colorable.NewColorable(w)
    46  	}
    47  	//////
    48  
    49  	doneWriting := make(chan bool)
    50  	result := make(chan []string)
    51  
    52  	go captureOutputAsyncronously(doneWriting, result, r)
    53  
    54  	block()
    55  	w.Close()
    56  	doneWriting <- true
    57  	return <-result
    58  }
    59  
    60  /*
    61   The reason we're doing is that you can't write an infinite amount of bytes into a pipe.
    62   On some platforms, the limit is fairly high; on other platforms, the limit is infuriatingly small
    63   (looking at you, Windows). To counteract this, we need to read in a goroutine from one end of
    64   the pipe and return the result across a channel.
    65  */
    66  func captureOutputAsyncronously(doneWriting <-chan bool, result chan<- []string, reader io.Reader) {
    67  	var readingString string
    68  
    69  	for {
    70  		var buf bytes.Buffer
    71  		io.Copy(&buf, reader)
    72  		readingString += buf.String()
    73  
    74  		_, ok := <-doneWriting
    75  		if ok {
    76  			// there is no guarantee that the writer did not
    77  			// write more in between the read above and reading from this channel
    78  			// so we absolute must read once more if we want all the bytes
    79  			var buf bytes.Buffer
    80  			io.Copy(&buf, reader)
    81  			readingString += buf.String()
    82  			break
    83  		}
    84  	}
    85  
    86  	result <- strings.Split(readingString, "\n")
    87  }