github.com/grafana/pyroscope@v1.18.0/pkg/test/capture.go (about)

     1  package test
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"os"
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  type CapturedOutput struct {
    14  	stdoutBuf bytes.Buffer
    15  	stderrBuf bytes.Buffer
    16  
    17  	wg                         sync.WaitGroup
    18  	stdoutReader, stdoutWriter *os.File
    19  	stderrReader, stderrWriter *os.File
    20  }
    21  
    22  // CaptureOutput replaces os.Stdout and os.Stderr with new pipes, that will
    23  // write output to buffers. Buffers are accessible by calling Done on returned
    24  // struct.
    25  //
    26  // os.Stdout and os.Stderr must be reverted to previous values manually.
    27  func CaptureOutput(t *testing.T) *CapturedOutput {
    28  	stdoutR, stdoutW, err := os.Pipe()
    29  	require.NoError(t, err)
    30  
    31  	stderrR, stderrW, err := os.Pipe()
    32  	require.NoError(t, err)
    33  
    34  	os.Stdout = stdoutW
    35  	os.Stderr = stderrW
    36  
    37  	co := &CapturedOutput{
    38  		stdoutReader: stdoutR,
    39  		stdoutWriter: stdoutW,
    40  		stderrReader: stderrR,
    41  		stderrWriter: stderrW,
    42  	}
    43  	co.wg.Add(1)
    44  	go func() {
    45  		defer co.wg.Done()
    46  		_, _ = io.Copy(&co.stdoutBuf, stdoutR)
    47  	}()
    48  
    49  	co.wg.Add(1)
    50  	go func() {
    51  		defer co.wg.Done()
    52  		_, _ = io.Copy(&co.stderrBuf, stderrR)
    53  	}()
    54  
    55  	return co
    56  }
    57  
    58  // Done waits until all captured output has been written to buffers,
    59  // and then returns the buffers.
    60  func (co *CapturedOutput) Done() (stdout string, stderr string) {
    61  	// we need to close writers for readers to stop
    62  	_ = co.stdoutWriter.Close()
    63  	_ = co.stderrWriter.Close()
    64  
    65  	co.wg.Wait()
    66  
    67  	return co.stdoutBuf.String(), co.stderrBuf.String()
    68  }