Source file src/pkg/testing/iotest/reader.go
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 returns a Reader that returns the final 41 // error with the last data read, instead of by itself with 42 // zero bytes of data. 43 func DataErrReader(r io.Reader) io.Reader { return &dataErrReader{r, nil, make([]byte, 1024)} } 44 45 type dataErrReader struct { 46 r io.Reader 47 unread []byte 48 data []byte 49 } 50 51 func (r *dataErrReader) Read(p []byte) (n int, err error) { 52 // loop because first call needs two reads: 53 // one to get data and a second to look for an error. 54 for { 55 if len(r.unread) == 0 { 56 n1, err1 := r.r.Read(r.data) 57 r.unread = r.data[0:n1] 58 err = err1 59 } 60 if n > 0 || err != nil { 61 break 62 } 63 n = copy(p, r.unread) 64 r.unread = r.unread[n:] 65 } 66 return 67 } 68 69 var ErrTimeout = errors.New("timeout") 70 71 // TimeoutReader returns ErrTimeout on the second read 72 // with no data. Subsequent calls to read succeed. 73 func TimeoutReader(r io.Reader) io.Reader { return &timeoutReader{r, 0} } 74 75 type timeoutReader struct { 76 r io.Reader 77 count int 78 } 79 80 func (r *timeoutReader) Read(p []byte) (int, error) { 81 r.count++ 82 if r.count == 2 { 83 return 0, ErrTimeout 84 } 85 return r.r.Read(p) 86 }