Add (Un)Marshal functions to stats.TagSet

stats/system_tag.go

@@ -23,6 +23,7 @@ package stats
 import (
 	"bytes"
 	"encoding/json"
+	"sort"
 	"strings"
 )
 
@@ -32,12 +33,57 @@ import (
 type SystemTagSet uint32
 
 // TagSet is a string to bool map (for lookup efficiency) that is used to keep track
-// which system tags and non-system tags should be included with with metrics.
+// of which system tags and non-system tags to include.
 type TagSet map[string]bool
 
-//nolint: golint
+// UnmarshalText converts the tag list to TagSet.
+func (i *TagSet) UnmarshalText(data []byte) error {
+	list := bytes.Split(data, []byte(","))
+	if *i == nil {
+		*i = make(TagSet, len(list))
+	}
+
+	for _, key := range list {
+		key := strings.TrimSpace(string(key))
+		if key == "" {
+			continue
+		}
+		(*i)[key] = true
+	}
+
+	return nil
+}
+
+// MarshalJSON converts the TagSet to a list (JS array).
+func (i *TagSet) MarshalJSON() ([]byte, error) {
+	var tags []string
+	if *i != nil {
+		tags = make([]string, 0, len(*i))
+		for tag := range *i {
+			tags = append(tags, tag)
+		}
+		sort.Strings(tags)
+	}
+
+	return json.Marshal(tags)
+}
+
+// UnmarshalJSON converts the tag list back to expected tag set.
+func (i *TagSet) UnmarshalJSON(data []byte) error {
+	var tags []string
+	if err := json.Unmarshal(data, &tags); err != nil {
+		return err
+	}
+	*i = make(TagSet, len(tags))
+	for _, tag := range tags {
+		(*i)[tag] = true
+	}
+
+	return nil
+}
+
+// Default system tags includes all of the system tags emitted with metrics by default.
 const (
-	// Default system tags includes all of the system tags emitted with metrics by default.
 	TagProto SystemTagSet = 1 << iota
 	TagSubproto
 	TagStatus
@@ -131,6 +177,8 @@ func (i *SystemTagSet) MarshalJSON() ([]byte, error) {
 			tags = append(tags, tag.String())
 		}
 	}
+	sort.Strings(tags)
+
 	return json.Marshal(tags)
 }
 
@@ -149,7 +197,7 @@ func (i *SystemTagSet) UnmarshalJSON(data []byte) error {
 
 // UnmarshalText converts the tag list to SystemTagSet.
 func (i *SystemTagSet) UnmarshalText(data []byte) error {
-	var list = bytes.Split(data, []byte(","))
+	list := bytes.Split(data, []byte(","))
 
 	for _, key := range list {
 		key := strings.TrimSpace(string(key))

stats/system_tag_test.go

@@ -29,11 +29,12 @@ import (
 )
 
 func TestSystemTagSetMarshalJSON(t *testing.T) {
-	var tests = []struct {
+	tests := []struct {
 		tagset   SystemTagSet
 		expected string
 	}{
 		{TagIP, `["ip"]`},
+		{TagIP | TagProto | TagGroup, `["group","ip","proto"]`},
 		{0, `null`},
 	}
 
@@ -46,7 +47,7 @@ func TestSystemTagSetMarshalJSON(t *testing.T) {
 }
 
 func TestSystemTagSet_UnmarshalJSON(t *testing.T) {
-	var tests = []struct {
+	tests := []struct {
 		tags []byte
 		sets []SystemTagSet
 	}{
@@ -64,7 +65,7 @@ func TestSystemTagSet_UnmarshalJSON(t *testing.T) {
 }
 
 func TestSystemTagSetTextUnmarshal(t *testing.T) {
-	var testMatrix = map[string]SystemTagSet{
+	testMatrix := map[string]SystemTagSet{
 		"":                      0,
 		"ip":                    TagIP,
 		"ip,proto":              TagIP | TagProto,
@@ -74,7 +75,61 @@ func TestSystemTagSetTextUnmarshal(t *testing.T) {
 	}
 
 	for input, expected := range testMatrix {
-		var set = new(SystemTagSet)
+		set := new(SystemTagSet)
+		err := set.UnmarshalText([]byte(input))
+		require.NoError(t, err)
+		require.Equal(t, expected, *set)
+	}
+}
+
+func TestTagSetMarshalJSON(t *testing.T) {
+	tests := []struct {
+		tagset   TagSet
+		expected string
+	}{
+		{tagset: TagSet{"ip": true, "proto": true, "group": true, "custom": true}, expected: `["custom","group","ip","proto"]`},
+		{tagset: TagSet{}, expected: `[]`},
+	}
+
+	for _, tc := range tests {
+		ts := &tc.tagset
+		got, err := json.Marshal(ts)
+		require.Nil(t, err)
+		require.Equal(t, tc.expected, string(got))
+	}
+}
+
+func TestTagSet_UnmarshalJSON(t *testing.T) {
+	tests := []struct {
+		tags []byte
+		sets TagSet
+	}{
+		{[]byte(`[]`), TagSet{}},
+		{[]byte(`["ip","custom", "proto"]`), TagSet{"ip": true, "proto": true, "custom": true}},
+	}
+
+	for _, tc := range tests {
+		ts := new(TagSet)
+		require.Nil(t, json.Unmarshal(tc.tags, ts))
+		for tag := range tc.sets {
+			assert.True(t, (*ts)[tag])
+		}
+	}
+}
+
+func TestTagSetTextUnmarshal(t *testing.T) {
+	testMatrix := map[string]TagSet{
+		"":                           make(TagSet),
+		"ip":                         {"ip": true},
+		"ip,proto":                   {"ip": true, "proto": true},
+		"   ip  ,  proto  ":          {"ip": true, "proto": true},
+		"   ip  ,   ,  proto  ":      {"ip": true, "proto": true},
+		"   ip  ,,  proto  ,,":       {"ip": true, "proto": true},
+		"   ip  ,custom,  proto  ,,": {"ip": true, "custom": true, "proto": true},
+	}
+
+	for input, expected := range testMatrix {
+		set := new(TagSet)
 		err := set.UnmarshalText([]byte(input))
 		require.NoError(t, err)
 		require.Equal(t, expected, *set)