Refactor template logic (#233)

Takashi Kusumi · web-flow · commit 371daf15cc58 · 2023-02-26T08:45:12.000+09:00
* Remove unnecessary no-color template

* Move template logic to function

* Add unit test for template logic
diff --git a/cmd/cmd.go b/cmd/cmd.go
@@ -234,74 +234,9 @@ func (o *options) sternConfig() (*stern.Config, error) {
 		return nil, errors.New("color should be one of 'always', 'never', or 'auto'")
 	}
 
-	t := o.template
-	if t == "" {
-		switch o.output {
-		case "default":
-			if color.NoColor {
-				t = "{{.PodName}} {{.ContainerName}} {{.Message}}"
-				if o.allNamespaces || len(o.namespaces) > 1 {
-					t = fmt.Sprintf("{{.Namespace}} %s", t)
-				}
-			} else {
-				t = "{{color .PodColor .PodName}} {{color .ContainerColor .ContainerName}} {{.Message}}"
-				if o.allNamespaces || len(o.namespaces) > 1 {
-					t = fmt.Sprintf("{{color .PodColor .Namespace}} %s", t)
-				}
-
-			}
-		case "raw":
-			t = "{{.Message}}"
-		case "json":
-			t = "{{json .}}"
-		case "extjson":
-			t = "\"pod\": \"{{color .PodColor .PodName}}\", \"container\": \"{{color .ContainerColor .ContainerName}}\", \"message\": {{extjson .Message}}"
-			if o.allNamespaces {
-				t = fmt.Sprintf("\"namespace\": \"{{color .PodColor .Namespace}}\", %s", t)
-			}
-			t = fmt.Sprintf("{%s}", t)
-		case "ppextjson":
-			t = "  \"pod\": \"{{color .PodColor .PodName}}\",\n  \"container\": \"{{color .ContainerColor .ContainerName}}\",\n  \"message\": {{extjson .Message}}"
-			if o.allNamespaces {
-				t = fmt.Sprintf("  \"namespace\": \"{{color .PodColor .Namespace}}\",\n%s", t)
-			}
-			t = fmt.Sprintf("{\n%s\n}", t)
-		}
-		t += "\n"
-	}
-
-	funs := map[string]interface{}{
-		"json": func(in interface{}) (string, error) {
-			b, err := json.Marshal(in)
-			if err != nil {
-				return "", err
-			}
-			return string(b), nil
-		},
-		"parseJSON": func(text string) (map[string]interface{}, error) {
-			obj := make(map[string]interface{})
-			if err := json.Unmarshal([]byte(text), &obj); err != nil {
-				return obj, err
-			}
-			return obj, nil
-		},
-		"extjson": func(in string) (string, error) {
-			if json.Valid([]byte(in)) {
-				return strings.TrimSuffix(in, "\n"), nil
-			}
-			b, err := json.Marshal(in)
-			if err != nil {
-				return "", err
-			}
-			return strings.TrimSuffix(string(b), "\n"), nil
-		},
-		"color": func(color color.Color, text string) string {
-			return color.SprintFunc()(text)
-		},
-	}
-	template, err := template.New("log").Funcs(funs).Parse(t)
+	template, err := o.generateTemplate()
 	if err != nil {
-		return nil, errors.Wrap(err, "unable to parse template")
+		return nil, err
 	}
 
 	namespaces := []string{}
@@ -402,6 +337,71 @@ func (o *options) AddFlags(fs *pflag.FlagSet) {
 	fs.BoolVarP(&o.version, "version", "v", o.version, "Print the version and exit.")
 }
 
+func (o *options) generateTemplate() (*template.Template, error) {
+	t := o.template
+	if t == "" {
+		switch o.output {
+		case "default":
+			t = "{{color .PodColor .PodName}} {{color .ContainerColor .ContainerName}} {{.Message}}"
+			if o.allNamespaces || len(o.namespaces) > 1 {
+				t = fmt.Sprintf("{{color .PodColor .Namespace}} %s", t)
+			}
+		case "raw":
+			t = "{{.Message}}"
+		case "json":
+			t = "{{json .}}"
+		case "extjson":
+			t = "\"pod\": \"{{color .PodColor .PodName}}\", \"container\": \"{{color .ContainerColor .ContainerName}}\", \"message\": {{extjson .Message}}"
+			if o.allNamespaces {
+				t = fmt.Sprintf("\"namespace\": \"{{color .PodColor .Namespace}}\", %s", t)
+			}
+			t = fmt.Sprintf("{%s}", t)
+		case "ppextjson":
+			t = "  \"pod\": \"{{color .PodColor .PodName}}\",\n  \"container\": \"{{color .ContainerColor .ContainerName}}\",\n  \"message\": {{extjson .Message}}"
+			if o.allNamespaces {
+				t = fmt.Sprintf("  \"namespace\": \"{{color .PodColor .Namespace}}\",\n%s", t)
+			}
+			t = fmt.Sprintf("{\n%s\n}", t)
+		}
+		t += "\n"
+	}
+
+	funs := map[string]interface{}{
+		"json": func(in interface{}) (string, error) {
+			b, err := json.Marshal(in)
+			if err != nil {
+				return "", err
+			}
+			return string(b), nil
+		},
+		"parseJSON": func(text string) (map[string]interface{}, error) {
+			obj := make(map[string]interface{})
+			if err := json.Unmarshal([]byte(text), &obj); err != nil {
+				return obj, err
+			}
+			return obj, nil
+		},
+		"extjson": func(in string) (string, error) {
+			if json.Valid([]byte(in)) {
+				return strings.TrimSuffix(in, "\n"), nil
+			}
+			b, err := json.Marshal(in)
+			if err != nil {
+				return "", err
+			}
+			return strings.TrimSuffix(string(b), "\n"), nil
+		},
+		"color": func(color color.Color, text string) string {
+			return color.SprintFunc()(text)
+		},
+	}
+	template, err := template.New("log").Funcs(funs).Parse(t)
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to parse template")
+	}
+	return template, err
+}
+
 func NewSternCmd(stream genericclioptions.IOStreams) (*cobra.Command, error) {
 	o := NewOptions(stream)
 
diff --git a/cmd/cmd_test.go b/cmd/cmd_test.go
@@ -1,9 +1,12 @@
 package cmd
 
 import (
+	"bytes"
 	"strings"
 	"testing"
 
+	"github.com/fatih/color"
+	"github.com/stern/stern/stern"
 	"k8s.io/cli-runtime/pkg/genericclioptions"
 )
 
@@ -126,3 +129,189 @@ func TestOptionsValidate(t *testing.T) {
 		})
 	}
 }
+
+func TestOptionsGenerateTemplate(t *testing.T) {
+	t.Setenv("NO_COLOR", "1")
+	streams := genericclioptions.NewTestIOStreamsDiscard()
+
+	tests := []struct {
+		name      string
+		o         *options
+		message   string
+		want      string
+		wantError bool
+	}{
+		{
+			"output=default",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "default"
+
+				return o
+			}(),
+			"default message",
+			"pod1 container1 default message\n",
+			false,
+		},
+		{
+			"output=default+allNamespaces",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "default"
+				o.allNamespaces = true
+
+				return o
+			}(),
+			"default message",
+			"ns1 pod1 container1 default message\n",
+			false,
+		},
+		{
+			"output=raw",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "raw"
+
+				return o
+			}(),
+			"raw message",
+			"raw message\n",
+			false,
+		},
+		{
+			"output=json",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "json"
+
+				return o
+			}(),
+			"json message",
+			`{"message":"json message","nodeName":"node1","namespace":"ns1","podName":"pod1","containerName":"container1"}
+`,
+			false,
+		},
+		{
+			"output=extjson",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "extjson"
+
+				return o
+			}(),
+			`{"msg":"extjson message"}`,
+			`{"pod": "pod1", "container": "container1", "message": {"msg":"extjson message"}}
+`,
+			false,
+		},
+		{
+			"output=extjson+allNamespaces",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "extjson"
+				o.allNamespaces = true
+
+				return o
+			}(),
+			`{"msg":"extjson message"}`,
+			`{"namespace": "ns1", "pod": "pod1", "container": "container1", "message": {"msg":"extjson message"}}
+`,
+			false,
+		},
+		{
+			"output=ppextjson",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "ppextjson"
+
+				return o
+			}(),
+			`{"msg":"ppextjson message"}`,
+			`{
+  "pod": "pod1",
+  "container": "container1",
+  "message": {"msg":"ppextjson message"}
+}
+`,
+			false,
+		},
+		{
+			"output=ppextjson+allNamespaces",
+			func() *options {
+				o := NewOptions(streams)
+				o.output = "ppextjson"
+				o.allNamespaces = true
+
+				return o
+			}(),
+			`{"msg":"ppextjson message"}`,
+			`{
+  "namespace": "ns1",
+  "pod": "pod1",
+  "container": "container1",
+  "message": {"msg":"ppextjson message"}
+}
+`,
+			false,
+		},
+		{
+			"template",
+			func() *options {
+				o := NewOptions(streams)
+				o.template = "Message={{.Message}} NodeName={{.NodeName}} Namespace={{.Namespace}} PodName={{.PodName}} ContainerName={{.ContainerName}}"
+
+				return o
+			}(),
+			"template message", // no new line
+			"Message=template message NodeName=node1 Namespace=ns1 PodName=pod1 ContainerName=container1",
+			false,
+		},
+		{
+			"invalid template",
+			func() *options {
+				o := NewOptions(streams)
+				o.template = "{{invalid"
+
+				return o
+			}(),
+			"template message",
+			"",
+			true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			log := stern.Log{
+				Message:        tt.message,
+				NodeName:       "node1",
+				Namespace:      "ns1",
+				PodName:        "pod1",
+				ContainerName:  "container1",
+				PodColor:       color.New(color.FgRed),
+				ContainerColor: color.New(color.FgBlue),
+			}
+			tmpl, err := tt.o.generateTemplate()
+
+			if tt.wantError {
+				if err == nil {
+					t.Errorf("expected error, but got no error")
+				}
+				return
+			}
+			if err != nil {
+				t.Errorf("unexpected error: %v", err)
+				return
+			}
+
+			var buf bytes.Buffer
+			if err := tmpl.Execute(&buf, log); err != nil {
+				t.Errorf("unexpected error: %v", err)
+				return
+			}
+			if want, got := tt.want, buf.String(); want != got {
+				t.Errorf("want %v, but got %v", want, got)
+			}
+		})
+	}
+}
