slog: handle removed built-in keys

Fix commonHandler so it works correctly when all built-in keys are
removed and there are preformatted attributes.
In that case, we need to add a separator between the
preformatted attributes and the rest.

Also, generalize the test for built-in handlers.  (Some features won't
prove useful until scope is introduced in a later CL).

Change-Id: If1fdeae34cc59da26f3d7a9bb1bb38c3ad14d366
Reviewed-on: https://go-review.googlesource.com/c/exp/+/438996
Run-TryBot: Jonathan Amsterdam <jba@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
diff --git a/slog/handler.go b/slog/handler.go
index f0c1f9e..6059ea1 100644
--- a/slog/handler.go
+++ b/slog/handler.go
@@ -230,6 +230,7 @@
 	if len(h.preformattedAttrs) > 0 {
 		state.buf.WriteString(state.sep)
 		state.buf.Write(h.preformattedAttrs)
+		state.sep = h.attrSep()
 	}
 	// Attrs in Record
 	r.Attrs(func(a Attr) {
@@ -246,6 +247,14 @@
 	return err
 }
 
+// attrSep returns the separator between attributes.
+func (h *commonHandler) attrSep() string {
+	if h.json {
+		return ","
+	}
+	return " "
+}
+
 // handleState holds state for a single call to commonHandler.handle.
 // The initial value of sep determines whether to emit a separator
 // before the next key, after which it stays true.
@@ -280,11 +289,10 @@
 	s.appendString(key)
 	if s.h.json {
 		s.buf.WriteByte(':')
-		s.sep = ","
 	} else {
 		s.buf.WriteByte('=')
-		s.sep = " "
 	}
+	s.sep = s.h.attrSep()
 }
 
 func (s *handleState) appendSource(file string, line int) {
diff --git a/slog/handler_test.go b/slog/handler_test.go
index 9d325e3..95fe219 100644
--- a/slog/handler_test.go
+++ b/slog/handler_test.go
@@ -58,7 +58,23 @@
 
 // Verify the common parts of TextHandler and JSONHandler.
 func TestJSONAndTextHandlers(t *testing.T) {
-	removeAttr := func(a Attr) Attr { return Attr{} }
+
+	// ReplaceAttr functions
+
+	// remove all Attrs
+	removeAll := func(a Attr) Attr { return Attr{} }
+
+	// remove the given keys
+	removeKeys := func(keys ...string) func(a Attr) Attr {
+		return func(a Attr) Attr {
+			for _, k := range keys {
+				if a.Key == k {
+					return Attr{}
+				}
+			}
+			return a
+		}
+	}
 
 	attrs := []Attr{String("a", "one"), Int("b", 2), Any("", "ignore me")}
 	preAttrs := []Attr{Int("pre", 3), String("x", "y")}
@@ -66,6 +82,7 @@
 	for _, test := range []struct {
 		name     string
 		replace  func(Attr) Attr
+		with     func(Handler) Handler
 		preAttrs []Attr
 		attrs    []Attr
 		wantText string
@@ -86,13 +103,14 @@
 		},
 		{
 			name:     "remove all",
-			replace:  removeAttr,
+			replace:  removeAll,
 			attrs:    attrs,
 			wantText: "",
 			wantJSON: `{}`,
 		},
 		{
 			name:     "preformatted",
+			with:     func(h Handler) Handler { return h.With(preAttrs) },
 			preAttrs: preAttrs,
 			attrs:    attrs,
 			wantText: "time=2000-01-02T03:04:05.000Z level=INFO msg=message pre=3 x=y a=one b=2",
@@ -101,6 +119,7 @@
 		{
 			name:     "preformatted cap keys",
 			replace:  upperCaseKey,
+			with:     func(h Handler) Handler { return h.With(preAttrs) },
 			preAttrs: preAttrs,
 			attrs:    attrs,
 			wantText: "TIME=2000-01-02T03:04:05.000Z LEVEL=INFO MSG=message PRE=3 X=y A=one B=2",
@@ -108,12 +127,28 @@
 		},
 		{
 			name:     "preformatted remove all",
-			replace:  removeAttr,
+			replace:  removeAll,
+			with:     func(h Handler) Handler { return h.With(preAttrs) },
 			preAttrs: preAttrs,
 			attrs:    attrs,
 			wantText: "",
 			wantJSON: "{}",
 		},
+		{
+			name:     "remove built-in",
+			replace:  removeKeys(timeKey, levelKey, messageKey),
+			attrs:    attrs,
+			wantText: "a=one b=2",
+			wantJSON: `{"a":"one","b":2}`,
+		},
+		{
+			name:     "preformatted remove built-in",
+			replace:  removeKeys(timeKey, levelKey, messageKey),
+			with:     func(h Handler) Handler { return h.With(preAttrs) },
+			attrs:    attrs,
+			wantText: "pre=3 x=y a=one b=2",
+			wantJSON: `{"pre":3,"x":"y","a":"one","b":2}`,
+		},
 	} {
 		r := NewRecord(testTime, InfoLevel, "message", 1)
 		r.AddAttrs(test.attrs...)
@@ -129,7 +164,10 @@
 				{"json", opts.NewJSONHandler(&buf), test.wantJSON},
 			} {
 				t.Run(handler.name, func(t *testing.T) {
-					h := handler.h.With(test.preAttrs)
+					h := handler.h
+					if test.with != nil {
+						h = test.with(h)
+					}
 					buf.Reset()
 					if err := h.Handle(r); err != nil {
 						t.Fatal(err)