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