github.com/mook-as/cf-cli@v7.0.0-beta.28.0.20200120190804-b91c115fae48+incompatible/cf/util/testhelpers/io/io.go (about)

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