stdcopy_test.go 7.9 KB

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