1 // Copyright 2009 The Go Authors. All rights reserved. 2 // Use of this source code is governed by a BSD-style 3 // license that can be found in the LICENSE file. 4 5 // Package iotest implements Readers and Writers useful mainly for testing. 6 package iotest 7 8 import ( 9 "errors" 10 "io" 11 ) 12 13 // OneByteReader returns a Reader that implements 14 // each non-empty Read by reading one byte from r. 15 func OneByteReader(r io.Reader) io.Reader { return &oneByteReader{r} } 16 17 type oneByteReader struct { 18 r io.Reader 19 } 20 21 func (r *oneByteReader) Read(p []byte) (int, error) { 22 if len(p) == 0 { 23 return 0, nil 24 } 25 return r.r.Read(p[0:1]) 26 } 27 28 // HalfReader returns a Reader that implements Read 29 // by reading half as many requested bytes from r. 30 func HalfReader(r io.Reader) io.Reader { return &halfReader{r} } 31 32 type halfReader struct { 33 r io.Reader 34 } 35 36 func (r *halfReader) Read(p []byte) (int, error) { 37 return r.r.Read(p[0 : (len(p)+1)/2]) 38 } 39 40 // DataErrReader changes the way errors are handled by a Reader. Normally, a 41 // Reader returns an error (typically EOF) from the first Read call after the 42 // last piece of data is read. DataErrReader wraps a Reader and changes its 43 // behavior so the final error is returned along with the final data, instead 44 // of in the first call after the final data. 45 func DataErrReader(r io.Reader) io.Reader { return &dataErrReader{r, nil, make([]byte, 1024)} } 46 47 type dataErrReader struct { 48 r io.Reader 49 unread []byte 50 data []byte 51 } 52 53 func (r *dataErrReader) Read(p []byte) (n int, err error) { 54 // loop because first call needs two reads: 55 // one to get data and a second to look for an error. 56 for { 57 if len(r.unread) == 0 { 58 n1, err1 := r.r.Read(r.data) 59 r.unread = r.data[0:n1] 60 err = err1 61 } 62 if n > 0 || err != nil { 63 break 64 } 65 n = copy(p, r.unread) 66 r.unread = r.unread[n:] 67 } 68 return 69 } 70 71 var ErrTimeout = errors.New("timeout") 72 73 // TimeoutReader returns ErrTimeout on the second read 74 // with no data. Subsequent calls to read succeed. 75 func TimeoutReader(r io.Reader) io.Reader { return &timeoutReader{r, 0} } 76 77 type timeoutReader struct { 78 r io.Reader 79 count int 80 } 81 82 func (r *timeoutReader) Read(p []byte) (int, error) { 83 r.count++ 84 if r.count == 2 { 85 return 0, ErrTimeout 86 } 87 return r.r.Read(p) 88 } 89
