| /* |
| Package gbytes provides a buffer that supports incrementally detecting input. |
| |
| 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. |
| |
| Subsequent matches against the buffer will only operate against data that appears *after* the read cursor. |
| |
| 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 |
| access the entire buffer's contents with Contents(). |
| |
| */ |
| package gbytes |
| |
| import ( |
| "errors" |
| "fmt" |
| "io" |
| "regexp" |
| "sync" |
| "time" |
| ) |
| |
| /* |
| gbytes.Buffer implements an io.Writer and can be used with the gbytes.Say matcher. |
| |
| 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! |
| */ |
| type Buffer struct { |
| contents []byte |
| readCursor uint64 |
| lock *sync.Mutex |
| detectCloser chan interface{} |
| closed bool |
| } |
| |
| /* |
| NewBuffer returns a new gbytes.Buffer |
| */ |
| func NewBuffer() *Buffer { |
| return &Buffer{ |
| lock: &sync.Mutex{}, |
| } |
| } |
| |
| /* |
| BufferWithBytes returns a new gbytes.Buffer seeded with the passed in bytes |
| */ |
| func BufferWithBytes(bytes []byte) *Buffer { |
| return &Buffer{ |
| lock: &sync.Mutex{}, |
| contents: bytes, |
| } |
| } |
| |
| /* |
| BufferReader returns a new gbytes.Buffer that wraps a reader. The reader's contents are read into |
| the Buffer via io.Copy |
| */ |
| func BufferReader(reader io.Reader) *Buffer { |
| b := &Buffer{ |
| lock: &sync.Mutex{}, |
| } |
| |
| go func() { |
| io.Copy(b, reader) |
| b.Close() |
| }() |
| |
| return b |
| } |
| |
| /* |
| Write implements the io.Writer interface |
| */ |
| func (b *Buffer) Write(p []byte) (n int, err error) { |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| if b.closed { |
| return 0, errors.New("attempt to write to closed buffer") |
| } |
| |
| b.contents = append(b.contents, p...) |
| return len(p), nil |
| } |
| |
| /* |
| Read implements the io.Reader interface. It advances the |
| cursor as it reads. |
| |
| Returns an error if called after Close. |
| */ |
| func (b *Buffer) Read(d []byte) (int, error) { |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| if b.closed { |
| return 0, errors.New("attempt to read from closed buffer") |
| } |
| |
| if uint64(len(b.contents)) <= b.readCursor { |
| return 0, io.EOF |
| } |
| |
| n := copy(d, b.contents[b.readCursor:]) |
| b.readCursor += uint64(n) |
| |
| return n, nil |
| } |
| |
| /* |
| Close signifies that the buffer will no longer be written to |
| */ |
| func (b *Buffer) Close() error { |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| b.closed = true |
| |
| return nil |
| } |
| |
| /* |
| Closed returns true if the buffer has been closed |
| */ |
| func (b *Buffer) Closed() bool { |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| return b.closed |
| } |
| |
| /* |
| Contents returns all data ever written to the buffer. |
| */ |
| func (b *Buffer) Contents() []byte { |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| contents := make([]byte, len(b.contents)) |
| copy(contents, b.contents) |
| return contents |
| } |
| |
| /* |
| Detect takes a regular expression and returns a channel. |
| |
| The channel will receive true the first time data matching the regular expression is written to the buffer. |
| The channel is subsequently closed and the buffer's read-cursor is fast-forwarded to just after the matching region. |
| |
| 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 |
| be branch and handle different outputs written to the buffer. |
| |
| 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. |
| |
| You could do something like: |
| |
| select { |
| case <-buffer.Detect("You are not logged in"): |
| //log in |
| case <-buffer.Detect("Success"): |
| //carry on |
| case <-time.After(time.Second): |
| //welp |
| } |
| buffer.CancelDetects() |
| |
| 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. |
| |
| Finally, you can pass detect a format string followed by variadic arguments. This will construct the regexp using fmt.Sprintf. |
| */ |
| func (b *Buffer) Detect(desired string, args ...interface{}) chan bool { |
| formattedRegexp := desired |
| if len(args) > 0 { |
| formattedRegexp = fmt.Sprintf(desired, args...) |
| } |
| re := regexp.MustCompile(formattedRegexp) |
| |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| if b.detectCloser == nil { |
| b.detectCloser = make(chan interface{}) |
| } |
| |
| closer := b.detectCloser |
| response := make(chan bool) |
| go func() { |
| ticker := time.NewTicker(10 * time.Millisecond) |
| defer ticker.Stop() |
| defer close(response) |
| for { |
| select { |
| case <-ticker.C: |
| b.lock.Lock() |
| data, cursor := b.contents[b.readCursor:], b.readCursor |
| loc := re.FindIndex(data) |
| b.lock.Unlock() |
| |
| if loc != nil { |
| response <- true |
| b.lock.Lock() |
| newCursorPosition := cursor + uint64(loc[1]) |
| if newCursorPosition >= b.readCursor { |
| b.readCursor = newCursorPosition |
| } |
| b.lock.Unlock() |
| return |
| } |
| case <-closer: |
| return |
| } |
| } |
| }() |
| |
| return response |
| } |
| |
| /* |
| 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. |
| */ |
| func (b *Buffer) CancelDetects() { |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| close(b.detectCloser) |
| b.detectCloser = nil |
| } |
| |
| func (b *Buffer) didSay(re *regexp.Regexp) (bool, []byte) { |
| b.lock.Lock() |
| defer b.lock.Unlock() |
| |
| unreadBytes := b.contents[b.readCursor:] |
| copyOfUnreadBytes := make([]byte, len(unreadBytes)) |
| copy(copyOfUnreadBytes, unreadBytes) |
| |
| loc := re.FindIndex(unreadBytes) |
| |
| if loc != nil { |
| b.readCursor += uint64(loc[1]) |
| return true, copyOfUnreadBytes |
| } |
| return false, copyOfUnreadBytes |
| } |