github.com/jchengjr77/canaveral@v1.0.1-0.20200715160102-ea9245d1a2fb/lib/iotesting.go (about)

     1  //Package lib contains testing functions for canaveral
     2  // includes:
     3  // - capture output (from stdout)
     4  // - redirect stdout
     5  // - reset stdout
     6  package lib
     7  
     8  import (
     9  	"bytes"
    10  	"io"
    11  	"log"
    12  	"os"
    13  	"sync"
    14  )
    15  
    16  // CaptureOutput takes in a function and reads all print statements.
    17  // Code snippet taken from:
    18  // https://medium.com/@hau12a1/golang-capturing-log-println-and-fmt-println-output-770209c791b4
    19  // * tested
    20  func CaptureOutput(f func()) string {
    21  	reader, writer, err := os.Pipe()
    22  	if err != nil {
    23  		panic(err)
    24  	}
    25  	stdout := os.Stdout
    26  	stderr := os.Stderr
    27  	defer func() {
    28  		os.Stdout = stdout
    29  		os.Stderr = stderr
    30  		log.SetOutput(os.Stderr)
    31  	}()
    32  	os.Stdout = writer
    33  	os.Stderr = writer
    34  	log.SetOutput(writer)
    35  	out := make(chan string)
    36  	wg := new(sync.WaitGroup)
    37  	wg.Add(1)
    38  	go func() {
    39  		var buf bytes.Buffer
    40  		wg.Done()
    41  		io.Copy(&buf, reader)
    42  		out <- buf.String()
    43  	}()
    44  	wg.Wait()
    45  	f()
    46  	writer.Close()
    47  	return <-out
    48  }
    49  
    50  // RedirOut redirects standard out away, and returns original *file.
    51  // This is mainly so whilst testing, the console doesn't get flooded.
    52  // * tested
    53  func RedirOut() *os.File {
    54  	_, writer, _ := os.Pipe()
    55  	realStdout := os.Stdout
    56  	os.Stdout = writer // redirect output away
    57  	return realStdout
    58  }
    59  
    60  // ResetOut resets os.stdout to the original stdout
    61  // * tested
    62  func ResetOut(stdout *os.File) {
    63  	tempstdOut := os.Stdout
    64  	os.Stdout = stdout
    65  	tempstdOut.Close()
    66  }