vendor/github.com/onsi/gomega/gbytes/buffer.go

   1 /*
   2 Package gbytes provides a buffer that supports incrementally detecting input.
   3
   4 You use gbytes.Buffer with the gbytes.Say matcher.  When Say finds a match, it fastforwards the buffer's read cursor to the end of that match.
   5
   6 Subsequent matches against the buffer will only operate against data that appears *after* the read cursor.
   7
   8 The read cursor is an opaque implementation detail that you cannot access.  You should use the Say matcher to sift through the buffer.  You can always
   9 access the entire buffer's contents with Contents().
  10
  11 */
  12 package gbytes
  13
  14 import (
  15         "errors"
  16         "fmt"
  17         "io"
  18         "regexp"
  19         "sync"
  20         "time"
  21 )
  22
  23 /*
  24 gbytes.Buffer implements an io.Writer and can be used with the gbytes.Say matcher.
  25
  26 You should only use a gbytes.Buffer in test code.  It stores all writes in an in-memory buffer - behavior that is inappropriate for production code!
  27 */
  28 type Buffer struct {
  29         contents     []byte
  30         readCursor   uint64
  31         lock         *sync.Mutex
  32         detectCloser chan interface{}
  33         closed       bool
  34 }
  35
  36 /*
  37 NewBuffer returns a new gbytes.Buffer
  38 */
  39 func NewBuffer() *Buffer {
  40         return &Buffer{
  41                 lock: &sync.Mutex{},
  42         }
  43 }
  44
  45 /*
  46 BufferWithBytes returns a new gbytes.Buffer seeded with the passed in bytes
  47 */
  48 func BufferWithBytes(bytes []byte) *Buffer {
  49         return &Buffer{
  50                 lock:     &sync.Mutex{},
  51                 contents: bytes,
  52         }
  53 }
  54
  55 /*
  56 Write implements the io.Writer interface
  57 */
  58 func (b *Buffer) Write(p []byte) (n int, err error) {
  59         b.lock.Lock()
  60         defer b.lock.Unlock()
  61
  62         if b.closed {
  63                 return 0, errors.New("attempt to write to closed buffer")
  64         }
  65
  66         b.contents = append(b.contents, p...)
  67         return len(p), nil
  68 }
  69
  70 /*
  71 Read implements the io.Reader interface. It advances the
  72 cursor as it reads.
  73
  74 Returns an error if called after Close.
  75 */
  76 func (b *Buffer) Read(d []byte) (int, error) {
  77         b.lock.Lock()
  78         defer b.lock.Unlock()
  79
  80         if b.closed {
  81                 return 0, errors.New("attempt to read from closed buffer")
  82         }
  83
  84         if uint64(len(b.contents)) <= b.readCursor {
  85                 return 0, io.EOF
  86         }
  87
  88         n := copy(d, b.contents[b.readCursor:])
  89         b.readCursor += uint64(n)
  90
  91         return n, nil
  92 }
  93
  94 /*
  95 Close signifies that the buffer will no longer be written to
  96 */
  97 func (b *Buffer) Close() error {
  98         b.lock.Lock()
  99         defer b.lock.Unlock()
 100
 101         b.closed = true
 102
 103         return nil
 104 }
 105
 106 /*
 107 Closed returns true if the buffer has been closed
 108 */
 109 func (b *Buffer) Closed() bool {
 110         b.lock.Lock()
 111         defer b.lock.Unlock()
 112
 113         return b.closed
 114 }
 115
 116 /*
 117 Contents returns all data ever written to the buffer.
 118 */
 119 func (b *Buffer) Contents() []byte {
 120         b.lock.Lock()
 121         defer b.lock.Unlock()
 122
 123         contents := make([]byte, len(b.contents))
 124         copy(contents, b.contents)
 125         return contents
 126 }
 127
 128 /*
 129 Detect takes a regular expression and returns a channel.
 130
 131 The channel will receive true the first time data matching the regular expression is written to the buffer.
 132 The channel is subsequently closed and the buffer's read-cursor is fast-forwarded to just after the matching region.
 133
 134 You typically don't need to use Detect and should use the ghttp.Say matcher instead.  Detect is useful, however, in cases where your code must
 135 be branch and handle different outputs written to the buffer.
 136
 137 For example, consider a buffer hooked up to the stdout of a client library.  You may (or may not, depending on state outside of your control) need to authenticate the client library.
 138
 139 You could do something like:
 140
 141 select {
 142 case <-buffer.Detect("You are not logged in"):
 143         //log in
 144 case <-buffer.Detect("Success"):
 145         //carry on
 146 case <-time.After(time.Second):
 147         //welp
 148 }
 149 buffer.CancelDetects()
 150
 151 You should always call CancelDetects after using Detect.  This will close any channels that have not detected and clean up the goroutines that were spawned to support them.
 152
 153 Finally, you can pass detect a format string followed by variadic arguments.  This will construct the regexp using fmt.Sprintf.
 154 */
 155 func (b *Buffer) Detect(desired string, args ...interface{}) chan bool {
 156         formattedRegexp := desired
 157         if len(args) > 0 {
 158                 formattedRegexp = fmt.Sprintf(desired, args...)
 159         }
 160         re := regexp.MustCompile(formattedRegexp)
 161
 162         b.lock.Lock()
 163         defer b.lock.Unlock()
 164
 165         if b.detectCloser == nil {
 166                 b.detectCloser = make(chan interface{})
 167         }
 168
 169         closer := b.detectCloser
 170         response := make(chan bool)
 171         go func() {
 172                 ticker := time.NewTicker(10 * time.Millisecond)
 173                 defer ticker.Stop()
 174                 defer close(response)
 175                 for {
 176                         select {
 177                         case <-ticker.C:
 178                                 b.lock.Lock()
 179                                 data, cursor := b.contents[b.readCursor:], b.readCursor
 180                                 loc := re.FindIndex(data)
 181                                 b.lock.Unlock()
 182
 183                                 if loc != nil {
 184                                         response <- true
 185                                         b.lock.Lock()
 186                                         newCursorPosition := cursor + uint64(loc[1])
 187                                         if newCursorPosition >= b.readCursor {
 188                                                 b.readCursor = newCursorPosition
 189                                         }
 190                                         b.lock.Unlock()
 191                                         return
 192                                 }
 193                         case <-closer:
 194                                 return
 195                         }
 196                 }
 197         }()
 198
 199         return response
 200 }
 201
 202 /*
 203 CancelDetects cancels any pending detects and cleans up their goroutines.  You should always call this when you're done with a set of Detect channels.
 204 */
 205 func (b *Buffer) CancelDetects() {
 206         b.lock.Lock()
 207         defer b.lock.Unlock()
 208
 209         close(b.detectCloser)
 210         b.detectCloser = nil
 211 }
 212
 213 func (b *Buffer) didSay(re *regexp.Regexp) (bool, []byte) {
 214         b.lock.Lock()
 215         defer b.lock.Unlock()
 216
 217         unreadBytes := b.contents[b.readCursor:]
 218         copyOfUnreadBytes := make([]byte, len(unreadBytes))
 219         copy(copyOfUnreadBytes, unreadBytes)
 220
 221         loc := re.FindIndex(unreadBytes)
 222
 223         if loc != nil {
 224                 b.readCursor += uint64(loc[1])
 225                 return true, copyOfUnreadBytes
 226         } else {
 227                 return false, copyOfUnreadBytes
 228         }
 229 }