From b89b363a76007ab521339afef55ae751ec7f07f4 Mon Sep 17 00:00:00 2001
From: Ryan Moran <rmoran@vmware.com>
Date: Wed, 22 Jun 2022 15:28:09 -0700
Subject: [PATCH] Adds scribe.WithPrefix option for scribe.Writer

---
 scribe/writer.go      | 13 +++++++++++++
 scribe/writer_test.go | 35 +++++++++++++++++++++++++++++++----
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/scribe/writer.go b/scribe/writer.go
index c5465bf9..026ac271 100644
--- a/scribe/writer.go
+++ b/scribe/writer.go
@@ -27,12 +27,23 @@ func WithIndent(indent int) Option {
 	}
 }
 
+// WithPrefix takes a prefix string and returns an Option which can be passed
+// in while creating a new Writer to configure a prefix to be prepended to the
+// output of the Writer.
+func WithPrefix(prefix string) Option {
+	return func(l Writer) Writer {
+		l.prefix = prefix
+		return l
+	}
+}
+
 // A Writer conforms to the io.Writer interface and allows for configuration of
 // output from the writter such as the color or indentation through Options.
 type Writer struct {
 	writer    io.Writer
 	color     Color
 	indent    int
+	prefix    string
 	linestart bool
 }
 
@@ -72,6 +83,8 @@ func (w *Writer) Write(b []byte) (int, error) {
 	var indentedLines [][]byte
 	for index, line := range lines {
 		if !(index == 0 && !w.linestart) {
+			line = append([]byte(w.prefix), line...)
+
 			for i := 0; i < w.indent; i++ {
 				line = append([]byte("  "), line...)
 			}
diff --git a/scribe/writer_test.go b/scribe/writer_test.go
index 6826b0b2..2ca2b7e1 100644
--- a/scribe/writer_test.go
+++ b/scribe/writer_test.go
@@ -70,27 +70,54 @@ func testWriter(t *testing.T, context spec.G, it spec.S) {
 				})
 			})
 
+			context("when the writer has a prefix", func() {
+				it.Before(func() {
+					writer = scribe.NewWriter(buffer, scribe.WithPrefix("[some-prefix] "))
+				})
+
+				it("prints to the writer with the given prefix", func() {
+					_, err := writer.Write([]byte("some-text\nother-text"))
+					Expect(err).NotTo(HaveOccurred())
+					Expect(buffer.String()).To(Equal("[some-prefix] some-text\n[some-prefix] other-text"))
+				})
+
+				context("when sequential write inputs are not newline terminated", func() {
+					it("handles the indentation correctly", func() {
+						_, err := writer.Write([]byte("some-text"))
+						Expect(err).NotTo(HaveOccurred())
+
+						_, err = writer.Write([]byte(" followed by other-text\n"))
+						Expect(err).NotTo(HaveOccurred())
+
+						_, err = writer.Write([]byte("followed by\neven-more-text\n"))
+						Expect(err).NotTo(HaveOccurred())
+
+						Expect(buffer.String()).To(Equal("[some-prefix] some-text followed by other-text\n[some-prefix] followed by\n[some-prefix] even-more-text\n"))
+					})
+				})
+			})
+
 			context("when the writer has a return prefix", func() {
 				it.Before(func() {
-					writer = scribe.NewWriter(buffer, scribe.WithColor(scribe.RedColor), scribe.WithIndent(2))
+					writer = scribe.NewWriter(buffer, scribe.WithColor(scribe.RedColor), scribe.WithIndent(2), scribe.WithPrefix("[some-prefix] "))
 				})
 
 				it("prints to the writer with the correct indentation", func() {
 					_, err := writer.Write([]byte("\rsome-text"))
 					Expect(err).NotTo(HaveOccurred())
-					Expect(buffer.String()).To(Equal("\r\x1b[0;38;5;1m    some-text\x1b[0m"))
+					Expect(buffer.String()).To(Equal("\r\x1b[0;38;5;1m    [some-prefix] some-text\x1b[0m"))
 				})
 			})
 
 			context("when the writer has a newline suffix", func() {
 				it.Before(func() {
-					writer = scribe.NewWriter(buffer, scribe.WithColor(scribe.RedColor), scribe.WithIndent(2))
+					writer = scribe.NewWriter(buffer, scribe.WithColor(scribe.RedColor), scribe.WithIndent(2), scribe.WithPrefix("[some-prefix] "))
 				})
 
 				it("prints to the writer with the correct indentation", func() {
 					_, err := writer.Write([]byte("some-text\n"))
 					Expect(err).NotTo(HaveOccurred())
-					Expect(buffer.String()).To(Equal("\x1b[0;38;5;1m    some-text\x1b[0m\n"))
+					Expect(buffer.String()).To(Equal("\x1b[0;38;5;1m    [some-prefix] some-text\x1b[0m\n"))
 				})
 			})