goutils/testio/testio.go

// Package testio implements various io utility types. Included are
// BrokenWriter, which fails after writing a certain number of bytes;
// a BufCloser, which wraps a bytes.Buffer in a Close method; a
// BrokenReadWriter, which fails after writing a certain number of
// bytes and/or reading a certain number of bytes; a LoggingBuffer
// that logs all reads and writes; and a BufferConn, that is designed
// to simulate net.Conn.
package testio

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
)

// BrokenWriter implements an io.Writer that fails after a certain
// number of bytes. This can be used to simulate a network connection
// that breaks during write or a file on a filesystem that becomes
// full, for example. A BrokenWriter doesn't actually store any data.
type BrokenWriter struct {
	current, limit int
}

// NewBrokenWriter creates a new BrokenWriter that can store only
// limit bytes.
func NewBrokenWriter(limit int) *BrokenWriter {
	return &BrokenWriter{limit: limit}
}

// Write will write the byte slice to the BrokenWriter, failing if the
// maximum number of bytes has been reached.
func (w *BrokenWriter) Write(p []byte) (int, error) {
	if (len(p) + w.current) <= w.limit {
		w.current += len(p)
	} else {
		spill := (len(p) + w.current) - w.limit
		w.current = w.limit
		return len(p) - spill, errors.New("testio: write failed")
	}

	return len(p), nil
}

// Extend increases the byte limit to allow more data to be written.
func (w *BrokenWriter) Extend(n int) {
	w.limit += n
}

// Reset clears the limit and bytes in the BrokenWriter. Extend needs
// to be called to allow data to be written.
func (w *BrokenWriter) Reset() {
	w.limit = 0
	w.current = 0
}

// Close is provided to satisfy the Closer interface.
func (w *BrokenWriter) Close() error {
	w.Reset()
	return nil
}

// SilentBrokenWriter implements an io.Writer that fails after a
// certain number of bytes. However, this failure is silent: it just
// reports fewer bytes written than p.  It doesn't actually store any
// data, and is used to verify that io.Writer implementations properly
// return errors on short writes.
type SilentBrokenWriter struct {
	current, limit int
}

// NewSilentBrokenWriter creates a new SilentBrokenWriter that can store only
// limit bytes.
func NewSilentBrokenWriter(limit int) *SilentBrokenWriter {
	return &SilentBrokenWriter{limit: limit}
}

// Write will write the byte slice to the SilentBrokenWriter, failing if the
// maximum number of bytes has been reached.
func (w *SilentBrokenWriter) Write(p []byte) (int, error) {
	if (len(p) + w.current) <= w.limit {
		w.current += len(p)
	} else {
		spill := (len(p) + w.current) - w.limit
		w.current = w.limit
		return len(p) - spill, nil
	}

	return len(p), nil
}

// Extend increases the byte limit to allow more data to be written.
func (w *SilentBrokenWriter) Extend(n int) {
	w.limit += n
}

// Reset clears the limit and bytes in the SilentBrokenWriter. Extend needs
// to be called to allow data to be written.
func (w *SilentBrokenWriter) Reset() {
	w.limit = 0
	w.current = 0
}

// Close is provided to satisfy the Closer interface.
func (w *SilentBrokenWriter) Close() error {
	w.Reset()
	return nil
}

// BrokenReadWriter implements a broken reader and writer, backed by a
// bytes.Buffer.
type BrokenReadWriter struct {
	rlimit, wlimit int
	buf            *bytes.Buffer
}

// NewBrokenReadWriter initialises a new BrokerReadWriter with an empty
// reader and the specified limits.
func NewBrokenReadWriter(wlimit, rlimit int) *BrokenReadWriter {
	return &BrokenReadWriter{
		wlimit: wlimit,
		rlimit: rlimit,
		buf:    &bytes.Buffer{},
	}
}

// Write satisfies the Writer interface.
func (brw *BrokenReadWriter) Write(p []byte) (int, error) {
	if (len(p) + brw.buf.Len()) > brw.wlimit {
		remain := brw.wlimit - brw.buf.Len()
		if remain > 0 {
			brw.buf.Write(p[:remain])
			return remain, errors.New("testio: write failed")
		}
		return 0, errors.New("testio: write failed")
	}
	return brw.buf.Write(p)
}

// Read satisfies the Reader interface.
func (brw *BrokenReadWriter) Read(p []byte) (int, error) {
	remain := brw.rlimit - brw.buf.Len()
	if len(p) > remain {
		tmp := make([]byte, len(p)-remain)
		n, err := brw.buf.Read(tmp)
		if err == nil {
			err = io.EOF
		}
		copy(p, tmp)
		return n, err
	}
	return brw.buf.Read(p)
}

// Extend increases the BrokenReadWriter limit.
func (brw *BrokenReadWriter) Extend(w, r int) {
	brw.rlimit += r
	brw.wlimit += w
}

// Reset clears the internal buffer. It retains its original limit.
func (brw *BrokenReadWriter) Reset() {
	brw.buf.Reset()
}

// BufCloser is a buffer wrapped with a Close method.
type BufCloser struct {
	buf *bytes.Buffer
}

// Write writes the data to the BufCloser.
func (buf *BufCloser) Write(p []byte) (int, error) {
	return buf.buf.Write(p)
}

// Read reads data from the BufCloser.
func (buf *BufCloser) Read(p []byte) (int, error) {
	return buf.buf.Read(p)
}

// Close is a stub function to satisfy the io.Closer interface.
func (buf *BufCloser) Close() error {
	return nil
}

// Reset clears the internal buffer.
func (buf *BufCloser) Reset() {
	buf.buf.Reset()
}

// Bytes returns the contents of the buffer as a byte slice.
func (buf *BufCloser) Bytes() []byte {
	return buf.buf.Bytes()
}

// Len returns the length of the buffer.
func (buf *BufCloser) Len() int {
	return buf.buf.Len()
}

// NewBufCloser creates and initializes a new BufCloser using buf as
// its initial contents. It is intended to prepare a BufCloser to read
// existing data. It can also be used to size the internal buffer for
// writing. To do that, buf should have the desired capacity but a
// length of zero.
func NewBufCloser(buf []byte) *BufCloser {
	bc := new(BufCloser)
	bc.buf = bytes.NewBuffer(buf)
	return bc
}

// NewBufCloserString creates and initializes a new Buffer using
// string s as its initial contents. It is intended to prepare a
// buffer to read an existing string.
func NewBufCloserString(s string) *BufCloser {
	buf := new(BufCloser)
	buf.buf = bytes.NewBufferString(s)
	return buf
}

// A LoggingBuffer is an io.ReadWriter that prints the hex value of
// the data for all reads and writes.
type LoggingBuffer struct {
	rw   io.ReadWriter
	w    io.Writer
	name string
}

// NewLoggingBuffer creates a logging buffer from an existing
// io.ReadWriter. By default, it will log to standard error.
func NewLoggingBuffer(rw io.ReadWriter) *LoggingBuffer {
	return &LoggingBuffer{
		rw: rw,
		w:  os.Stderr,
	}
}

// LogTo sets the io.Writer that the buffer will write logs to.
func (lb *LoggingBuffer) LogTo(w io.Writer) {
	lb.w = w
}

// SetName gives a name to the logging buffer to help distinguish
// output from this buffer.
func (lb *LoggingBuffer) SetName(name string) {
	lb.name = name
}

// Write writes the data to the logging buffer and writes the data to
// the logging writer.
func (lb *LoggingBuffer) Write(p []byte) (int, error) {
	if lb.name != "" {
		fmt.Fprintf(lb.w, "[%s] ", lb.name)
	}

	fmt.Fprintf(lb.w, "[WRITE] %x\n", p)
	return lb.rw.Write(p)
}

// Read reads the data from the logging buffer and writes the data to
// the logging writer.
func (lb *LoggingBuffer) Read(p []byte) (int, error) {
	n, err := lb.rw.Read(p)
	if err != nil {
		return n, err
	}
	if lb.name != "" {
		fmt.Fprintf(lb.w, "[%s] ", lb.name)
	}

	fmt.Fprintf(lb.w, "[READ] %x\n", p)
	return n, err
}

// BufferConn is a type that can be used to simulate network
// connections between a "client" (the code that uses the BufferConn)
// and some simulated "peer". Writes go to a "client" buffer, which is
// used to record the data sent by the caller, which may be read with
// ReadPeer. The peer's responses may be simulated by calling
// WritePeer; when the client reads from the BufferConn, they will see
// this data.
type BufferConn struct {
	client, peer *bytes.Buffer
}

// NewBufferConn initialises a new simulated network connection.
func NewBufferConn() *BufferConn {
	return &BufferConn{
		client: &bytes.Buffer{},
		peer:   &bytes.Buffer{},
	}
}

// Write writes to the client buffer.
func (bc *BufferConn) Write(p []byte) (int, error) {
	return bc.client.Write(p)
}

// Read reads from the peer buffer.
func (bc *BufferConn) Read(p []byte) (int, error) {
	return bc.peer.Read(p)
}

// WritePeer writes data to the peer buffer.
func (bc *BufferConn) WritePeer(p []byte) (int, error) {
	return bc.peer.Write(p)
}

// ReadClient reads data from the client buffer.
func (bc *BufferConn) ReadClient(p []byte) (int, error) {
	return bc.client.Read(p)
}

// Close is a dummy operation that allows the BufferConn to be used as
// an io.Closer.
func (bc *BufferConn) Close() error {
	return nil
}

// BrokenCloser is a BufCloser that fails to close.
type BrokenCloser struct {
	buf *bytes.Buffer
}

// Write writes the data to the BrokenCloser.
func (buf *BrokenCloser) Write(p []byte) (int, error) {
	return buf.buf.Write(p)
}

// Read reads data from the BrokenCloser.
func (buf *BrokenCloser) Read(p []byte) (int, error) {
	return buf.buf.Read(p)
}

// Close is a stub function to satisfy the io.Closer interface.
func (buf *BrokenCloser) Close() error {
	return errors.New("testio: broken closer is broken")
}

// Reset clears the internal buffer.
func (buf *BrokenCloser) Reset() {
	buf.buf.Reset()
}

// Bytes returns the contents of the buffer as a byte slice.
func (buf *BrokenCloser) Bytes() []byte {
	return buf.buf.Bytes()
}

// NewBrokenCloser creates and initializes a new BrokenCloser using buf as
// its initial contents. It is intended to prepare a BrokenCloser to read
// existing data. It can also be used to size the internal buffer for
// writing. To do that, buf should have the desired capacity but a
// length of zero.
func NewBrokenCloser(buf []byte) *BrokenCloser {
	bc := new(BrokenCloser)
	bc.buf = bytes.NewBuffer(buf)
	return bc
}

// NewBrokenCloserString creates and initializes a new Buffer using
// string s as its initial contents. It is intended to prepare a
// buffer to read an existing string.
func NewBrokenCloserString(s string) *BrokenCloser {
	buf := new(BrokenCloser)
	buf.buf = bytes.NewBufferString(s)
	return buf
}