stream.go

Documentation: github.com/linkerd/linkerd2/testutil

     1  package testutil
     2  
     3  import (
     4  	"bufio"
     5  	"fmt"
     6  	"io"
     7  	"os/exec"
     8  	"strings"
     9  	"time"
    10  )
    11  
    12  // Stream provides the ability of read the output of an executing process while
    13  // it is still running
    14  type Stream struct {
    15  	cmd *exec.Cmd
    16  	out io.ReadCloser
    17  }
    18  
    19  // Stop closes the stream and kills the process
    20  func (s *Stream) Stop() {
    21  	s.out.Close()
    22  	s.cmd.Process.Kill()
    23  }
    24  
    25  // ReadUntil reads from the process output until specified number of lines has
    26  // been reached, or until a timeout
    27  func (s *Stream) ReadUntil(lineCount int, timeout time.Duration) ([]string, error) {
    28  	output := make([]string, 0)
    29  	lines := make(chan string)
    30  	timeoutAfter := time.NewTimer(timeout)
    31  	defer timeoutAfter.Stop()
    32  	scanner := bufio.NewScanner(s.out)
    33  	stopSignal := false
    34  
    35  	go func() {
    36  		for scanner.Scan() {
    37  			lines <- scanner.Text()
    38  
    39  			if stopSignal {
    40  				close(lines)
    41  				return
    42  			}
    43  		}
    44  	}()
    45  
    46  	for {
    47  		select {
    48  		case <-timeoutAfter.C:
    49  			stopSignal = true
    50  			return output, fmt.Errorf("cmd [%s] Timed out trying to read %d lines", strings.Join(s.cmd.Args, " "), lineCount)
    51  		case line := <-lines:
    52  			output = append(output, line)
    53  			if len(output) >= lineCount {
    54  				stopSignal = true
    55  				return output, nil
    56  			}
    57  		}
    58  	}
    59  }
    60
View as plain text