1  package testutils
  2  
  3  import (
  4  	"bufio"
  5  	"bytes"
  6  	"fmt"
  7  	"log"
  8  	"os"
  9  	"strings"
 10  	"testing"
 11  )
 12  
 13  // A LogHarness runs sets of assertions against the log output of a function. Assertions are grouped
 14  // into sequences of messages that are expected to be found in the log output. Calling one of the Expect
 15  // methods on the harness adds an expectation to the default sequence of messages. Additional sequences
 16  // can be created by calling NewSequence.
 17  type LogHarness struct {
 18  	buf       bytes.Buffer
 19  	sequences []*Sequence
 20  }
 21  
 22  type Expectation interface {
 23  	IsMatch(line string) bool
 24  	String() string
 25  }
 26  
 27  // Run executes the function f and captures any output written using Go's standard log. Each sequence
 28  // of expected messages is then
 29  func (h *LogHarness) Run(t *testing.T, f func()) {
 30  	// Capture raw log output
 31  	fl := log.Flags()
 32  	log.SetFlags(0)
 33  	log.SetOutput(&h.buf)
 34  	f()
 35  	log.SetFlags(fl)
 36  	log.SetOutput(os.Stderr)
 37  
 38  	for _, seq := range h.sequences {
 39  		seq.Assert(t, bufio.NewScanner(bytes.NewReader(h.buf.Bytes())))
 40  	}
 41  }
 42  
 43  // Expect adds an expectation to the default sequence that the log contains a line equal to s
 44  func (h *LogHarness) Expect(s string) {
 45  	if len(h.sequences) == 0 {
 46  		h.sequences = append(h.sequences, &Sequence{name: ""})
 47  	}
 48  	h.sequences[0].Expect(s)
 49  }
 50  
 51  // ExpectPrefix adds an to the default sequence expectation that the log contains a line starting with s
 52  func (h *LogHarness) ExpectPrefix(s string) {
 53  	if len(h.sequences) == 0 {
 54  		h.sequences = append(h.sequences, &Sequence{name: ""})
 55  	}
 56  	h.sequences[0].ExpectPrefix(s)
 57  }
 58  
 59  // NewSequence creates a new sequence of expected log messages
 60  func (h *LogHarness) NewSequence(name string) *Sequence {
 61  	seq := &Sequence{name: name}
 62  	h.sequences = append(h.sequences, seq)
 63  	return seq
 64  }
 65  
 66  type prefix string
 67  
 68  func (p prefix) IsMatch(line string) bool {
 69  	return strings.HasPrefix(line, string(p))
 70  }
 71  
 72  func (p prefix) String() string {
 73  	return fmt.Sprintf("prefix %q", string(p))
 74  }
 75  
 76  type text string
 77  
 78  func (t text) IsMatch(line string) bool {
 79  	return line == string(t)
 80  }
 81  
 82  func (t text) String() string {
 83  	return fmt.Sprintf("text %q", string(t))
 84  }
 85  
 86  type Sequence struct {
 87  	name string
 88  	exp  []Expectation
 89  }
 90  
 91  func (seq *Sequence) Assert(t *testing.T, s *bufio.Scanner) {
 92  	var tag string
 93  	if seq.name != "" {
 94  		tag = fmt.Sprintf("[%s] ", seq.name)
 95  	}
 96  	// Match raw log lines against expectations
 97  exploop:
 98  	for _, e := range seq.exp {
 99  		for s.Scan() {
100  			if e.IsMatch(s.Text()) {
101  				t.Logf("%ssaw: %s", tag, s.Text())
102  				continue exploop
103  			}
104  		}
105  		if s.Err() == nil {
106  			t.Errorf("%sdid not see expected %s", tag, e.String())
107  			return
108  		}
109  	}
110  }
111  
112  // Expect adds an expectation that the log contains a line equal to s
113  func (seq *Sequence) Expect(s string) {
114  	seq.exp = append(seq.exp, text(s))
115  }
116  
117  // ExpectPrefix adds an expectation that the log contains a line starting with s
118  func (seq *Sequence) ExpectPrefix(s string) {
119  	seq.exp = append(seq.exp, prefix(s))
120  }