hclwrite: fix TokensForTraversal handling of index steps

2 files changed, 48 insertions, 3 deletions

diff --git a/hclwrite/generate.go b/hclwrite/generate.go
index 289a30d..4d439ac 100644
--- a/hclwrite/generate.go
+++ b/hclwrite/generate.go
@@ -159,12 +159,12 @@ func appendTokensForValue(val cty.Value, toks Tokens) Tokens {
 
 func appendTokensForTraversal(traversal hcl.Traversal, toks Tokens) Tokens {
 	for _, step := range traversal {
-		appendTokensForTraversalStep(step, toks)
+		toks = appendTokensForTraversalStep(step, toks)
 	}
 	return toks
 }
 
-func appendTokensForTraversalStep(step hcl.Traverser, toks Tokens) {
+func appendTokensForTraversalStep(step hcl.Traverser, toks Tokens) Tokens {
 	switch ts := step.(type) {
 	case hcl.TraverseRoot:
 		toks = append(toks, &Token{
@@ -188,7 +188,7 @@ func appendTokensForTraversalStep(step hcl.Traverser, toks Tokens) {
 			Type:  hclsyntax.TokenOBrack,
 			Bytes: []byte{'['},
 		})
-		appendTokensForValue(ts.Key, toks)
+		toks = appendTokensForValue(ts.Key, toks)
 		toks = append(toks, &Token{
 			Type:  hclsyntax.TokenCBrack,
 			Bytes: []byte{']'},
@@ -196,6 +196,8 @@ func appendTokensForTraversalStep(step hcl.Traverser, toks Tokens) {
 	default:
 		panic(fmt.Sprintf("unsupported traversal step type %T", step))
 	}
+
+	return toks
 }
 
 func escapeQuotedStringLit(s string) []byte {
diff --git a/hclwrite/generate_test.go b/hclwrite/generate_test.go
index 7c437e2..112ff82 100644
--- a/hclwrite/generate_test.go
+++ b/hclwrite/generate_test.go
@@ -6,6 +6,7 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/hashicorp/hcl/v2"
 	"github.com/hashicorp/hcl/v2/hclsyntax"
 	"github.com/zclconf/go-cty/cty"
 )
@@ -472,3 +473,45 @@ func TestTokensForValue(t *testing.T) {
 		})
 	}
 }
+
+func TestTokensForTraversal(t *testing.T) {
+	tests := []struct {
+		Val  hcl.Traversal
+		Want Tokens
+	}{
+		{
+			hcl.Traversal{
+				hcl.TraverseRoot{Name: "root"},
+				hcl.TraverseAttr{Name: "attr"},
+				hcl.TraverseIndex{Key: cty.StringVal("index")},
+			},
+			Tokens{
+				{Type: hclsyntax.TokenIdent, Bytes: []byte("root")},
+				{Type: hclsyntax.TokenDot, Bytes: []byte(".")},
+				{Type: hclsyntax.TokenIdent, Bytes: []byte("attr")},
+				{Type: hclsyntax.TokenOBrack, Bytes: []byte{'['}},
+				{Type: hclsyntax.TokenOQuote, Bytes: []byte(`"`)},
+				{Type: hclsyntax.TokenQuotedLit, Bytes: []byte("index")},
+				{Type: hclsyntax.TokenCQuote, Bytes: []byte(`"`)},
+				{Type: hclsyntax.TokenCBrack, Bytes: []byte{']'}},
+			},
+		},
+	}
+
+	for _, test := range tests {
+		got := TokensForTraversal(test.Val)
+
+		if !cmp.Equal(got, test.Want) {
+			diff := cmp.Diff(got, test.Want, cmp.Comparer(func(a, b []byte) bool {
+				return bytes.Equal(a, b)
+			}))
+			var gotBuf, wantBuf bytes.Buffer
+			got.WriteTo(&gotBuf)
+			test.Want.WriteTo(&wantBuf)
+			t.Errorf(
+				"wrong result\nvalue: %#v\ngot:   %s\nwant:  %s\ndiff:  %s",
+				test.Val, gotBuf.String(), wantBuf.String(), diff,
+			)
+		}
+	}
+}