goutils/testio/testio.go

262 lines
6.9 KiB
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
}
// 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()
}
// 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
}