github.com/jmigpin/editor@v1.6.0/util/testutil/collectlog.go (about)

     1  package testutil
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"os"
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/jmigpin/editor/util/iout"
    11  )
    12  
    13  func CollectLog(t *testing.T, fn func() error) ([]byte, []byte, error) {
    14  	return CollectLog2(t.Logf, fn)
    15  }
    16  func CollectLog2(logf func(string, ...any), fn func() error) ([]byte, []byte, error) {
    17  	// keep for later restoration
    18  	orig1, orig2 := os.Stdout, os.Stderr
    19  	defer func() { // restore
    20  		os.Stdout, os.Stderr = orig1, orig2
    21  	}()
    22  
    23  	// build pipes to catch stdout/stderr for ouput comparison
    24  	stdoutBuf, stderrBuf := &bytes.Buffer{}, &bytes.Buffer{}
    25  	pr1, pw1, err1 := os.Pipe()
    26  	if err1 != nil {
    27  		return nil, nil, err1
    28  	}
    29  	pr2, pw2, err2 := os.Pipe()
    30  	if err2 != nil {
    31  		return nil, nil, err2
    32  	}
    33  	os.Stdout, os.Stderr = pw1, pw2
    34  
    35  	// setup logger
    36  	logWriter := func(wr io.Writer, buf *bytes.Buffer) io.Writer {
    37  		return iout.FnWriter(func(b []byte) (int, error) {
    38  			// commented: prints many lines without log prefix
    39  			//logf("%s", b)
    40  
    41  			// ensure a call to logf() sends a complete line
    42  			k := 0
    43  			for i, c := range b {
    44  				if c == '\n' {
    45  					b2 := append(buf.Bytes(), b[k:i]...)
    46  					logf("%s", b2)
    47  					k = i + 1
    48  					buf.Reset()
    49  				}
    50  			}
    51  			if k < len(b) {
    52  				buf.Write(b[k:])
    53  			}
    54  			return len(b), nil
    55  		})
    56  	}
    57  	flushLogWriter := func(buf *bytes.Buffer) {
    58  		if len(buf.Bytes()) > 0 {
    59  			logf("%s", buf.Bytes())
    60  		}
    61  	}
    62  
    63  	// copy loop
    64  	copyLoops := &sync.WaitGroup{}
    65  	copyLoops.Add(1)
    66  	go func() {
    67  		defer copyLoops.Done()
    68  
    69  		buf := &bytes.Buffer{}
    70  		lw := logWriter(orig1, buf)
    71  
    72  		rd := io.Reader(pr1)
    73  		rd = io.TeeReader(rd, lw)
    74  		_, _ = io.Copy(stdoutBuf, rd)
    75  		pr1.Close()
    76  
    77  		flushLogWriter(buf) // final bytes
    78  	}()
    79  
    80  	// copy loop
    81  	copyLoops.Add(1)
    82  	go func() {
    83  		defer copyLoops.Done()
    84  
    85  		buf := &bytes.Buffer{}
    86  		lw := logWriter(orig2, buf)
    87  
    88  		rd := io.Reader(pr2)
    89  		rd = io.TeeReader(rd, lw)
    90  		_, _ = io.Copy(stderrBuf, rd)
    91  		pr2.Close()
    92  
    93  		flushLogWriter(buf) // final bytes
    94  	}()
    95  
    96  	// run
    97  	err := fn()
    98  
    99  	// cmd done, close writers
   100  	pw1.Close()
   101  	pw2.Close()
   102  	// wait for readers
   103  	copyLoops.Wait()
   104  
   105  	return stdoutBuf.Bytes(), stderrBuf.Bytes(), err
   106  }