From 3cc15b53e7698c20241900538160cfc2e9787536 Mon Sep 17 00:00:00 2001
From: Vincent Demeester <vincent@sbr.pm>
Date: Thu, 16 Apr 2015 14:17:23 +0200
Subject: [PATCH] Add some stdcopy_test (coverage)

Signed-off-by: Vincent Demeester <vincent@sbr.pm>
---
 stdcopy/stdcopy.go      |  4 ---
 stdcopy/stdcopy_test.go | 65 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/stdcopy/stdcopy.go b/stdcopy/stdcopy.go
index 4f78f20..dbb74e5 100644
--- a/stdcopy/stdcopy.go
+++ b/stdcopy/stdcopy.go
@@ -54,10 +54,6 @@ func (w *StdWriter) Write(buf []byte) (n int, err error) {
 // `t` indicates the id of the stream to encapsulate.
 // It can be stdcopy.Stdin, stdcopy.Stdout, stdcopy.Stderr.
 func NewStdWriter(w io.Writer, t StdType) *StdWriter {
-	if len(t) != StdWriterPrefixLen {
-		return nil
-	}
-
 	return &StdWriter{
 		Writer:  w,
 		prefix:  t,
diff --git a/stdcopy/stdcopy_test.go b/stdcopy/stdcopy_test.go
index 14e6ed3..a9fd73a 100644
--- a/stdcopy/stdcopy_test.go
+++ b/stdcopy/stdcopy_test.go
@@ -3,9 +3,74 @@ package stdcopy
 import (
 	"bytes"
 	"io/ioutil"
+	"strings"
 	"testing"
 )
 
+func TestNewStdWriter(t *testing.T) {
+	writer := NewStdWriter(ioutil.Discard, Stdout)
+	if writer == nil {
+		t.Fatalf("NewStdWriter with an invalid StdType should not return nil.")
+	}
+}
+
+func TestWriteWithUnitializedStdWriter(t *testing.T) {
+	writer := StdWriter{
+		Writer:  nil,
+		prefix:  Stdout,
+		sizeBuf: make([]byte, 4),
+	}
+	n, err := writer.Write([]byte("Something here"))
+	if n != 0 || err == nil {
+		t.Fatalf("Should fail when given an uncomplete or uninitialized StdWriter")
+	}
+}
+
+func TestWriteWithNilBytes(t *testing.T) {
+	writer := NewStdWriter(ioutil.Discard, Stdout)
+	n, err := writer.Write(nil)
+	if err != nil {
+		t.Fatalf("Shouldn't have fail when given no data")
+	}
+	if n > 0 {
+		t.Fatalf("Write should have written 0 byte, but has written %d", n)
+	}
+}
+
+func TestWrite(t *testing.T) {
+	writer := NewStdWriter(ioutil.Discard, Stdout)
+	data := []byte("Test StdWrite.Write")
+	n, err := writer.Write(data)
+	if err != nil {
+		t.Fatalf("Error while writing with StdWrite")
+	}
+	if n != len(data) {
+		t.Fatalf("Write should have writen %d byte but wrote %d.", len(data), n)
+	}
+}
+
+func TestStdCopyWithInvalidInputHeader(t *testing.T) {
+	dstOut := NewStdWriter(ioutil.Discard, Stdout)
+	dstErr := NewStdWriter(ioutil.Discard, Stderr)
+	src := strings.NewReader("Invalid input")
+	_, err := StdCopy(dstOut, dstErr, src)
+	if err == nil {
+		t.Fatal("StdCopy with invalid input header should fail.")
+	}
+}
+
+func TestStdCopyWithCorruptedPrefix(t *testing.T) {
+	data := []byte{0x01, 0x02, 0x03}
+	src := bytes.NewReader(data)
+	written, err := StdCopy(nil, nil, src)
+	if err != nil {
+		t.Fatalf("StdCopy should not return an error with corrupted prefix.")
+	}
+	if written != 0 {
+		t.Fatalf("StdCopy should have written 0, but has written %d", written)
+	}
+}
+
 func BenchmarkWrite(b *testing.B) {
 	w := NewStdWriter(ioutil.Discard, Stdout)
 	data := []byte("Test line for testing stdwriter performance\n")