From d5a94653da883ea59205751e421a81953988b216 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=9F=D0=B0=D0=B2=D0=B5=D0=BB?= <pashok9825@gmail.com>
Date: Wed, 8 Apr 2020 13:50:02 +0300
Subject: [PATCH] fixes in template (un)marshaling and tests

---
 v1/template.go      | 14 +++++++---
 v1/template_test.go | 65 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 4 deletions(-)
 create mode 100644 v1/template_test.go

diff --git a/v1/template.go b/v1/template.go
index d88db9d..d852047 100644
--- a/v1/template.go
+++ b/v1/template.go
@@ -11,7 +11,7 @@ const TemplateTypeText = "text"
 
 const (
 	// TemplateItemTypeText is a type for text chunk in template
-	TemplateItemTypeText = iota
+	TemplateItemTypeText uint8 = iota
 	// TemplateItemTypeVar is a type for variable in template
 	TemplateItemTypeVar
 )
@@ -52,12 +52,17 @@ type TemplateItem struct {
 	VarType string
 }
 
-// MarshalJSON will correctly marshal TemplateItem
-func (t *TemplateItem) MarshalJSON() ([]byte, error) {
+// MarshalJSON controls how TemplateItem will be marshaled into JSON
+func (t TemplateItem) MarshalJSON() ([]byte, error) {
 	switch t.Type {
 	case TemplateItemTypeText:
 		return json.Marshal(t.Text)
 	case TemplateItemTypeVar:
+		// {} case, fast output without marshaling
+		if t.VarType == "" || t.VarType == TemplateVarCustom {
+			return []byte("{}"), nil
+		}
+
 		return json.Marshal(map[string]interface{}{
 			"var": t.VarType,
 		})
@@ -87,11 +92,12 @@ func (t *TemplateItem) UnmarshalJSON(b []byte) error {
 		}
 
 		if varTypeCurr, ok := v["var"].(string); ok {
-			if _, ok := templateVarAssoc[t.VarType]; !ok {
+			if _, ok := templateVarAssoc[varTypeCurr]; !ok {
 				return fmt.Errorf("invalid placeholder var '%s'", varTypeCurr)
 			}
 
 			t.Type = TemplateItemTypeVar
+			t.VarType = varTypeCurr
 		} else {
 			return errors.New("invalid TemplateItem")
 		}
diff --git a/v1/template_test.go b/v1/template_test.go
new file mode 100644
index 0000000..c9d9f26
--- /dev/null
+++ b/v1/template_test.go
@@ -0,0 +1,65 @@
+package v1
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTemplateItem_MarshalJSON(t *testing.T) {
+	text := TemplateItem{
+		Type: TemplateItemTypeText,
+		Text: "text item",
+	}
+
+	variable := TemplateItem{
+		Type:    TemplateItemTypeVar,
+		VarType: TemplateVarFirstName,
+	}
+
+	emptyVariable := TemplateItem{
+		Type:    TemplateItemTypeVar,
+		VarType: "",
+	}
+
+	data, err := json.Marshal(text)
+	assert.NoError(t, err)
+	assert.Equal(t, "\""+text.Text+"\"", string(data))
+
+	data, err = json.Marshal(variable)
+	assert.NoError(t, err)
+	assert.Equal(t, `{"var":"first_name"}`, string(data))
+
+	data, err = json.Marshal(emptyVariable)
+	assert.NoError(t, err)
+	assert.Equal(t, "{}", string(data))
+}
+
+func TestTemplateItem_UnmarshalJSON(t *testing.T) {
+	var (
+		textResult          TemplateItem
+		variableResult      TemplateItem
+		emptyVariableResult TemplateItem
+	)
+
+	text := []byte("\"text block\"")
+	variable := []byte(`{"var":"first_name"}`)
+	emptyVariable := []byte("{}")
+
+	require.NoError(t, json.Unmarshal(text, &textResult))
+	require.NoError(t, json.Unmarshal(variable, &variableResult))
+	require.NoError(t, json.Unmarshal(emptyVariable, &emptyVariableResult))
+
+	assert.Equal(t, TemplateItemTypeText, textResult.Type)
+	assert.Equal(t, string(text)[1:11], textResult.Text)
+
+	assert.Equal(t, TemplateItemTypeVar, variableResult.Type)
+	assert.Equal(t, TemplateVarFirstName, variableResult.VarType)
+	assert.Empty(t, variableResult.Text)
+
+	assert.Equal(t, TemplateItemTypeVar, emptyVariableResult.Type)
+	assert.Equal(t, TemplateVarCustom, emptyVariableResult.VarType)
+	assert.Empty(t, emptyVariableResult.Text)
+}