buffer.go

Documentation: github.com/onsi/gomega/gbytes

     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  BufferReader returns a new gbytes.Buffer that wraps a reader.  The reader's contents are read into
    57  the Buffer via io.Copy
    58  */
    59  func BufferReader(reader io.Reader) *Buffer {
    60  	b := &Buffer{
    61  		lock: &sync.Mutex{},
    62  	}
    63  
    64  	go func() {
    65  		io.Copy(b, reader)
    66  		b.Close()
    67  	}()
    68  
    69  	return b
    70  }
    71  
    72  /*
    73  Write implements the io.Writer interface
    74  */
    75  func (b *Buffer) Write(p []byte) (n int, err error) {
    76  	b.lock.Lock()
    77  	defer b.lock.Unlock()
    78  
    79  	if b.closed {
    80  		return 0, errors.New("attempt to write to closed buffer")
    81  	}
    82  
    83  	b.contents = append(b.contents, p...)
    84  	return len(p), nil
    85  }
    86  
    87  /*
    88  Read implements the io.Reader interface. It advances the
    89  cursor as it reads.
    90  */
    91  func (b *Buffer) Read(d []byte) (int, error) {
    92  	b.lock.Lock()
    93  	defer b.lock.Unlock()
    94  
    95  	if uint64(len(b.contents)) <= b.readCursor {
    96  		return 0, io.EOF
    97  	}
    98  
    99  	n := copy(d, b.contents[b.readCursor:])
   100  	b.readCursor += uint64(n)
   101  
   102  	return n, nil
   103  }
   104  
   105  /*
   106  Clear clears out the buffer's contents
   107  */
   108  func (b *Buffer) Clear() error {
   109  	b.lock.Lock()
   110  	defer b.lock.Unlock()
   111  
   112  	if b.closed {
   113  		return errors.New("attempt to clear closed buffer")
   114  	}
   115  
   116  	b.contents = []byte{}
   117  	b.readCursor = 0
   118  	return nil
   119  }
   120  
   121  /*
   122  Close signifies that the buffer will no longer be written to
   123  */
   124  func (b *Buffer) Close() error {
   125  	b.lock.Lock()
   126  	defer b.lock.Unlock()
   127  
   128  	b.closed = true
   129  
   130  	return nil
   131  }
   132  
   133  /*
   134  Closed returns true if the buffer has been closed
   135  */
   136  func (b *Buffer) Closed() bool {
   137  	b.lock.Lock()
   138  	defer b.lock.Unlock()
   139  
   140  	return b.closed
   141  }
   142  
   143  /*
   144  Contents returns all data ever written to the buffer.
   145  */
   146  func (b *Buffer) Contents() []byte {
   147  	b.lock.Lock()
   148  	defer b.lock.Unlock()
   149  
   150  	contents := make([]byte, len(b.contents))
   151  	copy(contents, b.contents)
   152  	return contents
   153  }
   154  
   155  /*
   156  Detect takes a regular expression and returns a channel.
   157  
   158  The channel will receive true the first time data matching the regular expression is written to the buffer.
   159  The channel is subsequently closed and the buffer's read-cursor is fast-forwarded to just after the matching region.
   160  
   161  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
   162  be branch and handle different outputs written to the buffer.
   163  
   164  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.
   165  
   166  You could do something like:
   167  
   168  select {
   169  case <-buffer.Detect("You are not logged in"):
   170  	//log in
   171  case <-buffer.Detect("Success"):
   172  	//carry on
   173  case <-time.After(time.Second):
   174  	//welp
   175  }
   176  buffer.CancelDetects()
   177  
   178  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.
   179  
   180  Finally, you can pass detect a format string followed by variadic arguments.  This will construct the regexp using fmt.Sprintf.
   181  */
   182  func (b *Buffer) Detect(desired string, args ...interface{}) chan bool {
   183  	formattedRegexp := desired
   184  	if len(args) > 0 {
   185  		formattedRegexp = fmt.Sprintf(desired, args...)
   186  	}
   187  	re := regexp.MustCompile(formattedRegexp)
   188  
   189  	b.lock.Lock()
   190  	defer b.lock.Unlock()
   191  
   192  	if b.detectCloser == nil {
   193  		b.detectCloser = make(chan interface{})
   194  	}
   195  
   196  	closer := b.detectCloser
   197  	response := make(chan bool)
   198  	go func() {
   199  		ticker := time.NewTicker(10 * time.Millisecond)
   200  		defer ticker.Stop()
   201  		defer close(response)
   202  		for {
   203  			select {
   204  			case <-ticker.C:
   205  				b.lock.Lock()
   206  				data, cursor := b.contents[b.readCursor:], b.readCursor
   207  				loc := re.FindIndex(data)
   208  				b.lock.Unlock()
   209  
   210  				if loc != nil {
   211  					response <- true
   212  					b.lock.Lock()
   213  					newCursorPosition := cursor + uint64(loc[1])
   214  					if newCursorPosition >= b.readCursor {
   215  						b.readCursor = newCursorPosition
   216  					}
   217  					b.lock.Unlock()
   218  					return
   219  				}
   220  			case <-closer:
   221  				return
   222  			}
   223  		}
   224  	}()
   225  
   226  	return response
   227  }
   228  
   229  /*
   230  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.
   231  */
   232  func (b *Buffer) CancelDetects() {
   233  	b.lock.Lock()
   234  	defer b.lock.Unlock()
   235  
   236  	close(b.detectCloser)
   237  	b.detectCloser = nil
   238  }
   239  
   240  func (b *Buffer) didSay(re *regexp.Regexp) (bool, []byte) {
   241  	b.lock.Lock()
   242  	defer b.lock.Unlock()
   243  
   244  	unreadBytes := b.contents[b.readCursor:]
   245  	copyOfUnreadBytes := make([]byte, len(unreadBytes))
   246  	copy(copyOfUnreadBytes, unreadBytes)
   247  
   248  	loc := re.FindIndex(unreadBytes)
   249  
   250  	if loc != nil {
   251  		b.readCursor += uint64(loc[1])
   252  		return true, copyOfUnreadBytes
   253  	}
   254  	return false, copyOfUnreadBytes
   255  }
   256
View as plain text