pkg/stdcopy/stdcopy_test.go

   1 package stdcopy
   2
   3 import (
   4         "bytes"
   5         "errors"
   6         "io"
   7         "io/ioutil"
   8         "strings"
   9         "testing"
  10 )
  11
  12 func TestNewStdWriter(t *testing.T) {
  13         writer := NewStdWriter(ioutil.Discard, Stdout)
  14         if writer == nil {
  15                 t.Fatalf("NewStdWriter with an invalid StdType should not return nil.")
  16         }
  17 }
  18
  19 func TestWriteWithUnitializedStdWriter(t *testing.T) {
  20         writer := stdWriter{
  21                 Writer: nil,
  22                 prefix: byte(Stdout),
  23         }
  24         n, err := writer.Write([]byte("Something here"))
  25         if n != 0 || err == nil {
  26                 t.Fatalf("Should fail when given an uncomplete or uninitialized StdWriter")
  27         }
  28 }
  29
  30 func TestWriteWithNilBytes(t *testing.T) {
  31         writer := NewStdWriter(ioutil.Discard, Stdout)
  32         n, err := writer.Write(nil)
  33         if err != nil {
  34                 t.Fatalf("Shouldn't have fail when given no data")
  35         }
  36         if n > 0 {
  37                 t.Fatalf("Write should have written 0 byte, but has written %d", n)
  38         }
  39 }
  40
  41 func TestWrite(t *testing.T) {
  42         writer := NewStdWriter(ioutil.Discard, Stdout)
  43         data := []byte("Test StdWrite.Write")
  44         n, err := writer.Write(data)
  45         if err != nil {
  46                 t.Fatalf("Error while writing with StdWrite")
  47         }
  48         if n != len(data) {
  49                 t.Fatalf("Write should have written %d byte but wrote %d.", len(data), n)
  50         }
  51 }
  52
  53 type errWriter struct {
  54         n   int
  55         err error
  56 }
  57
  58 func (f *errWriter) Write(buf []byte) (int, error) {
  59         return f.n, f.err
  60 }
  61
  62 func TestWriteWithWriterError(t *testing.T) {
  63         expectedError := errors.New("expected")
  64         expectedReturnedBytes := 10
  65         writer := NewStdWriter(&errWriter{
  66                 n:   stdWriterPrefixLen + expectedReturnedBytes,
  67                 err: expectedError}, Stdout)
  68         data := []byte("This won't get written, sigh")
  69         n, err := writer.Write(data)
  70         if err != expectedError {
  71                 t.Fatalf("Didn't get expected error.")
  72         }
  73         if n != expectedReturnedBytes {
  74                 t.Fatalf("Didn't get expected written bytes %d, got %d.",
  75                         expectedReturnedBytes, n)
  76         }
  77 }
  78
  79 func TestWriteDoesNotReturnNegativeWrittenBytes(t *testing.T) {
  80         writer := NewStdWriter(&errWriter{n: -1}, Stdout)
  81         data := []byte("This won't get written, sigh")
  82         actual, _ := writer.Write(data)
  83         if actual != 0 {
  84                 t.Fatalf("Expected returned written bytes equal to 0, got %d", actual)
  85         }
  86 }
  87
  88 func getSrcBuffer(stdOutBytes, stdErrBytes []byte) (buffer *bytes.Buffer, err error) {
  89         buffer = new(bytes.Buffer)
  90         dstOut := NewStdWriter(buffer, Stdout)
  91         _, err = dstOut.Write(stdOutBytes)
  92         if err != nil {
  93                 return
  94         }
  95         dstErr := NewStdWriter(buffer, Stderr)
  96         _, err = dstErr.Write(stdErrBytes)
  97         return
  98 }
  99
 100 func TestStdCopyWriteAndRead(t *testing.T) {
 101         stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
 102         stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
 103         buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
 104         if err != nil {
 105                 t.Fatal(err)
 106         }
 107         written, err := StdCopy(ioutil.Discard, ioutil.Discard, buffer)
 108         if err != nil {
 109                 t.Fatal(err)
 110         }
 111         expectedTotalWritten := len(stdOutBytes) + len(stdErrBytes)
 112         if written != int64(expectedTotalWritten) {
 113                 t.Fatalf("Expected to have total of %d bytes written, got %d", expectedTotalWritten, written)
 114         }
 115 }
 116
 117 type customReader struct {
 118         n            int
 119         err          error
 120         totalCalls   int
 121         correctCalls int
 122         src          *bytes.Buffer
 123 }
 124
 125 func (f *customReader) Read(buf []byte) (int, error) {
 126         f.totalCalls++
 127         if f.totalCalls <= f.correctCalls {
 128                 return f.src.Read(buf)
 129         }
 130         return f.n, f.err
 131 }
 132
 133 func TestStdCopyReturnsErrorReadingHeader(t *testing.T) {
 134         expectedError := errors.New("error")
 135         reader := &customReader{
 136                 err: expectedError}
 137         written, err := StdCopy(ioutil.Discard, ioutil.Discard, reader)
 138         if written != 0 {
 139                 t.Fatalf("Expected 0 bytes read, got %d", written)
 140         }
 141         if err != expectedError {
 142                 t.Fatalf("Didn't get expected error")
 143         }
 144 }
 145
 146 func TestStdCopyReturnsErrorReadingFrame(t *testing.T) {
 147         expectedError := errors.New("error")
 148         stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
 149         stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
 150         buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
 151         if err != nil {
 152                 t.Fatal(err)
 153         }
 154         reader := &customReader{
 155                 correctCalls: 1,
 156                 n:            stdWriterPrefixLen + 1,
 157                 err:          expectedError,
 158                 src:          buffer}
 159         written, err := StdCopy(ioutil.Discard, ioutil.Discard, reader)
 160         if written != 0 {
 161                 t.Fatalf("Expected 0 bytes read, got %d", written)
 162         }
 163         if err != expectedError {
 164                 t.Fatalf("Didn't get expected error")
 165         }
 166 }
 167
 168 func TestStdCopyDetectsCorruptedFrame(t *testing.T) {
 169         stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
 170         stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
 171         buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
 172         if err != nil {
 173                 t.Fatal(err)
 174         }
 175         reader := &customReader{
 176                 correctCalls: 1,
 177                 n:            stdWriterPrefixLen + 1,
 178                 err:          io.EOF,
 179                 src:          buffer}
 180         written, err := StdCopy(ioutil.Discard, ioutil.Discard, reader)
 181         if written != startingBufLen {
 182                 t.Fatalf("Expected %d bytes read, got %d", startingBufLen, written)
 183         }
 184         if err != nil {
 185                 t.Fatal("Didn't get nil error")
 186         }
 187 }
 188
 189 func TestStdCopyWithInvalidInputHeader(t *testing.T) {
 190         dstOut := NewStdWriter(ioutil.Discard, Stdout)
 191         dstErr := NewStdWriter(ioutil.Discard, Stderr)
 192         src := strings.NewReader("Invalid input")
 193         _, err := StdCopy(dstOut, dstErr, src)
 194         if err == nil {
 195                 t.Fatal("StdCopy with invalid input header should fail.")
 196         }
 197 }
 198
 199 func TestStdCopyWithCorruptedPrefix(t *testing.T) {
 200         data := []byte{0x01, 0x02, 0x03}
 201         src := bytes.NewReader(data)
 202         written, err := StdCopy(nil, nil, src)
 203         if err != nil {
 204                 t.Fatalf("StdCopy should not return an error with corrupted prefix.")
 205         }
 206         if written != 0 {
 207                 t.Fatalf("StdCopy should have written 0, but has written %d", written)
 208         }
 209 }
 210
 211 func TestStdCopyReturnsWriteErrors(t *testing.T) {
 212         stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
 213         stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
 214         buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
 215         if err != nil {
 216                 t.Fatal(err)
 217         }
 218         expectedError := errors.New("expected")
 219
 220         dstOut := &errWriter{err: expectedError}
 221
 222         written, err := StdCopy(dstOut, ioutil.Discard, buffer)
 223         if written != 0 {
 224                 t.Fatalf("StdCopy should have written 0, but has written %d", written)
 225         }
 226         if err != expectedError {
 227                 t.Fatalf("Didn't get expected error, got %v", err)
 228         }
 229 }
 230
 231 func TestStdCopyDetectsNotFullyWrittenFrames(t *testing.T) {
 232         stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
 233         stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
 234         buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
 235         if err != nil {
 236                 t.Fatal(err)
 237         }
 238         dstOut := &errWriter{n: startingBufLen - 10}
 239
 240         written, err := StdCopy(dstOut, ioutil.Discard, buffer)
 241         if written != 0 {
 242                 t.Fatalf("StdCopy should have return 0 written bytes, but returned %d", written)
 243         }
 244         if err != io.ErrShortWrite {
 245                 t.Fatalf("Didn't get expected io.ErrShortWrite error")
 246         }
 247 }
 248
 249 // TestStdCopyReturnsErrorFromSystem tests that StdCopy correctly returns an
 250 // error, when that error is muxed into the Systemerr stream.
 251 func TestStdCopyReturnsErrorFromSystem(t *testing.T) {
 252         // write in the basic messages, just so there's some fluff in there
 253         stdOutBytes := []byte(strings.Repeat("o", startingBufLen))
 254         stdErrBytes := []byte(strings.Repeat("e", startingBufLen))
 255         buffer, err := getSrcBuffer(stdOutBytes, stdErrBytes)
 256         if err != nil {
 257                 t.Fatal(err)
 258         }
 259         // add in an error message on the Systemerr stream
 260         systemErrBytes := []byte(strings.Repeat("S", startingBufLen))
 261         systemWriter := NewStdWriter(buffer, Systemerr)
 262         _, err = systemWriter.Write(systemErrBytes)
 263         if err != nil {
 264                 t.Fatal(err)
 265         }
 266
 267         // now copy and demux. we should expect an error containing the string we
 268         // wrote out
 269         _, err = StdCopy(ioutil.Discard, ioutil.Discard, buffer)
 270         if err == nil {
 271                 t.Fatal("expected error, got none")
 272         }
 273         if !strings.Contains(err.Error(), string(systemErrBytes)) {
 274                 t.Fatal("expected error to contain message")
 275         }
 276 }
 277
 278 func BenchmarkWrite(b *testing.B) {
 279         w := NewStdWriter(ioutil.Discard, Stdout)
 280         data := []byte("Test line for testing stdwriter performance\n")
 281         data = bytes.Repeat(data, 100)
 282         b.SetBytes(int64(len(data)))
 283         b.ResetTimer()
 284         for i := 0; i < b.N; i++ {
 285                 if _, err := w.Write(data); err != nil {
 286                         b.Fatal(err)
 287                 }
 288         }
 289 }