writer: simplify MarshalJSON + test extra fields roundtrip

gelf/writer.go

@@ -302,24 +302,25 @@ func (w *Writer) Write(p []byte) (n int, err error) {
 func (m *Message) MarshalJSON() ([]byte, error) {
  var err error
  var b, eb []byte
- if len(m.Extra) > 0 {
- if eb, err = json.Marshal(m.Extra); err != nil {
- return nil, err
- }
- }
+
  extra := m.Extra
- im := (*innerMessage)(m)
- im.Extra = nil
- if b, err = json.Marshal(im); err != nil {
- m.Extra = extra
- return b, err
+ b, err = json.Marshal((*innerMessage)(m))
+ m.Extra = extra
+ if err != nil {
+ return nil, err
  }
- if extra != nil && len(extra) > 0 && len(eb) > 0 {
- m.Extra = extra
- //fmt.Printf("b=%s\neb=%s\n", b, eb)
- return append(append(b[:len(b)-1], ','), eb[1:len(eb)]...), nil
+
+ if len(extra) == 0 {
+ return b, nil
  }
- return b, nil
+
+ if eb, err = json.Marshal(extra); err != nil {
+ return nil, err
+ }
+
+ // merge serialized message + serialized extra map
+ b[len(b)-1] = ','
+ return append(b, eb[1:len(eb)]...), nil
 }
 
 func (m *Message) UnmarshalJSON(data []byte) error {

gelf/writer_test.go

@@ -10,6 +10,7 @@ import (
  "fmt"
  "strings"
  "testing"
+ "time"
 )
 
 func TestNewWriter(t *testing.T) {
@@ -39,6 +40,25 @@ func sendAndRecv(msgData string, compress CompressType) (*Message, error) {
  return r.ReadMessage()
 }
 
+func sendAndRecvMsg(msg *Message, compress CompressType) (*Message, error) {
+ r, err := NewReader("127.0.0.1:0")
+ if err != nil {
+ return nil, fmt.Errorf("NewReader: %s", err)
+ }
+
+ w, err := NewWriter(r.Addr())
+ if err != nil {
+ return nil, fmt.Errorf("NewWriter: %s", err)
+ }
+ w.CompressionType = compress
+
+ if err = w.WriteMessage(msg); err != nil {
+ return nil, fmt.Errorf("w.Write: %s", err)
+ }
+
+ return r.ReadMessage()
+}
+
 // tests single-message (non-chunked) messages that are split over
 // multiple lines
 func TestWriteSmallMultiLine(t *testing.T) {
@@ -151,4 +171,58 @@ func TestWriteBigChunked(t *testing.T) {
  }
 }
 
-// tests single-message (non-chunked) messages that are a single line long
+// tests messages with extra data
+func TestExtraData(t *testing.T) {
+
+ // time.Now().Unix() seems fine, UnixNano() won't roundtrip
+ // through string -> float64 -> int64
+ extra := map[string]interface{}{"_a": 10 * time.Now().Unix(), "C": 9}
+
+ short := "quick"
+ full := short + "\nwith more detail"
+ m := Message{
+ Version: "1.0",
+ Host: "fake-host",
+ Short: string(short),
+ Full: string(full),
+ TimeUnix: time.Now().Unix(),
+ Level: 6, // info
+ Facility: "writer_test",
+ File: "writer_test.go",
+ Line: 186,
+ Extra: extra,
+ }
+
+ for _, i := range []CompressType{CompressGzip, CompressZlib} {
+ msg, err := sendAndRecvMsg(&m, i)
+ if err != nil {
+ t.Errorf("sendAndRecv: %s", err)
+ return
+ }
+
+ if msg.Short != short {
+ t.Errorf("msg.Short: expected %s, got %s", short, msg.Full)
+ return
+ }
+
+ if msg.Full != full {
+ t.Errorf("msg.Full: expected %s, got %s", full, msg.Full)
+ return
+ }
+
+ if msg.File != "writer_test.go" {
+ t.Errorf("msg.File: expected writer_test.go, got %s", msg.File)
+ return
+ }
+
+ if len(msg.Extra) != 1 {
+ t.Errorf("extra extra fields in %v", msg.Extra)
+ return
+ }
+
+ if int64(msg.Extra["_a"].(float64)) != extra["_a"].(int64) {
+ t.Errorf("_a didn't roundtrip (%v != %v)", int64(msg.Extra["_a"].(float64)), extra["_a"].(int64))
+ return
+ }
+ }
+}