123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245 |
- /*
- 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
- }
|