fix func logger and add test

Author: rdallman

api/agent/func_logger.go

@@ -111,18 +111,20 @@ func (li *lineWriter) Write(ogb []byte) (int, error) {
 
 	var n int
 	for {
-		// read the line and advance buffer past it
-		l, err := li.b.ReadBytes('\n')
-		if err != nil {
-			break // no more newlines in buffer (see ReadBytes contract)
+		b := li.b.Bytes()
+		i := bytes.IndexByte(b, '\n')
+		if i < 0 {
+			break // no more newlines in buffer
 		}
 
-		// write in the line
+		// write in this line and advance buffer past it
+		l := b[:i+1]
 		ns, err := li.w.Write(l)
 		n += ns
 		if err != nil {
 			return n, err
 		}
+		li.b.Next(len(l))
 	}
 
 	// technically we wrote all the bytes, so make things appear normal

api/agent/func_logger_test.go

@@ -0,0 +1,49 @@
+package agent
+
+import (
+	"bytes"
+	"fmt"
+	"testing"
+)
+
+type testSliceWriter struct {
+	b [][]byte
+}
+
+func (tsw *testSliceWriter) Write(p []byte) (n int, err error) {
+	l := make([]byte, len(p))
+	copy(l, p)
+	tsw.b = append(tsw.b, l)
+	return len(p), nil
+}
+
+func TestLineWriter(t *testing.T) {
+	var tsw testSliceWriter
+	lw := newLineWriter(&nopCloser{&tsw})
+
+	lineCount := 7
+	lw.Write([]byte("0 line\n1 line\n2 line\n\n4 line"))
+	lw.Write([]byte("+more\n5 line\n"))
+	lw.Write([]byte("6 line"))
+
+	lw.Close()
+
+	if len(tsw.b) != lineCount {
+		t.Errorf("Expected %v individual rows; got %v", lineCount, len(tsw.b))
+	}
+
+	for x := 0; x < len(tsw.b); x++ {
+		l := fmt.Sprintf("%v line\n", x)
+		if x == 3 {
+			if len(tsw.b[x]) != 1 {
+				t.Errorf("Expected slice with only newline; got %v", tsw.b[x])
+			}
+			continue
+		} else if x == 4 {
+			l = "4 line+more\n"
+		}
+		if !bytes.Equal(tsw.b[x], []byte(l)) {
+			t.Errorf("Expected slice %s equal to %s", []byte(l), tsw.b[x])
+		}
+	}
+}