fix: handle root variable ($) at different stage

mrjosh · May 20, 2024 · a26b767 · a26b767
1 parent da5d643
commit a26b767
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 18 deletions.
diff --git a/Makefile b/Makefile
@@ -68,7 +68,7 @@ test:
 	@$(GO) test ./... -v -race -tags=integration
 
 coverage:
-	@$(GO) test -coverprofile=.coverage -tags=integration ./internal/... && go tool cover -html=.coverage
+	@$(GO) test -coverprofile=.coverage -tags=integration -coverpkg=./internal/... ./internal/... && go tool cover -html=.coverage
 
 .PHONY: build-release
 build-release:

diff --git a/internal/lsp/symbol_table.go b/internal/lsp/symbol_table.go
@@ -45,7 +45,12 @@ func NewSymbolTable(ast *sitter.Tree, content []byte) *SymbolTable {
 }
 
 func (s *SymbolTable) AddTemplateContext(templateContext TemplateContext, pointRange sitter.Range) {
-	s.contexts[templateContext.Format()] = append(s.contexts[strings.Join(templateContext, ".")], pointRange)
+	if templateContext.IsVariable() && templateContext[0] == "$" {
+		// $ is a special variable that resolves to the root context
+		// we can just remove it from the template context
+		templateContext = templateContext.Tail()
+	}
+	s.contexts[templateContext.Format()] = append(s.contexts[templateContext.Format()], pointRange)
 	sliceCopy := make(TemplateContext, len(templateContext))
 	copy(sliceCopy, templateContext)
 	s.contextsReversed[pointRange] = sliceCopy

diff --git a/internal/lsp/symbol_table_template_context.go b/internal/lsp/symbol_table_template_context.go
@@ -60,7 +60,7 @@ func (v *TemplateContextVisitor) Enter(node *sitter.Node) {
 		v.symbolTable.AddTemplateContext(append(v.currentContext, content), GetRangeForNode(node.ChildByFieldName("name")))
 	case gotemplate.NodeTypeUnfinishedSelectorExpression:
 		operandNode := node.ChildByFieldName("operand")
-		if operandNode.Type() == gotemplate.NodeTypeVariable && operandNode.Content(v.content) == "$" {
+		if operandNode.Type() == gotemplate.NodeTypeVariable {
 			v.StashContext()
 		}
 		v.symbolTable.AddTemplateContext(append(getContextForSelectorExpression(operandNode, v.content), ""),
@@ -69,10 +69,8 @@ func (v *TemplateContextVisitor) Enter(node *sitter.Node) {
 		operandNode := node.ChildByFieldName("operand")
 		if operandNode.Type() == gotemplate.NodeTypeVariable {
 			v.StashContext()
-			if operandNode.Content(v.content) != "$" {
-				v.symbolTable.AddTemplateContext(append(v.currentContext, operandNode.Content(v.content)), GetRangeForNode(operandNode))
-				v.PushContext(operandNode.Content(v.content))
-			}
+			v.symbolTable.AddTemplateContext(append(v.currentContext, operandNode.Content(v.content)), GetRangeForNode(operandNode))
+			v.PushContext(operandNode.Content(v.content))
 		}
 	}
 }
@@ -82,9 +80,7 @@ func (v *TemplateContextVisitor) Exit(node *sitter.Node) {
 	case gotemplate.NodeTypeSelectorExpression, gotemplate.NodeTypeUnfinishedSelectorExpression:
 		operandNode := node.ChildByFieldName("operand")
 		if operandNode.Type() == gotemplate.NodeTypeVariable {
-			if operandNode.Content(v.content) != "$" {
-				v.PopContext()
-			}
+			v.PopContext()
 			v.RestoreStashedContext()
 		}
 	}
@@ -104,9 +100,6 @@ func (v *TemplateContextVisitor) EnterContextShift(node *sitter.Node, suffix str
 			s = s.AppendSuffix(suffix)
 			if s.IsVariable() {
 				v.StashContext()
-				if s[0] == "$" {
-					s = s.Tail()
-				}
 			}
 		}
 		v.PushContextMany(s)

diff --git a/internal/lsp/symbol_table_template_context_test.go b/internal/lsp/symbol_table_template_context_test.go
@@ -14,17 +14,41 @@ func TestGetContextForSelectorExpression(t *testing.T) {
 		expected    TemplateContext
 	}{
 		{
-			desc:        "Selects single selector expression correctly",
+			desc:        "Selects simple selector expression correctly",
 			template:    `{{ .Values.test }}`,
 			nodeContent: ".Values.test",
 			expected:    TemplateContext{"Values", "test"},
 		},
+		{
+			desc:        "Selects unfinished selector expression correctly",
+			template:    `{{ .Values.test. }}`,
+			nodeContent: ".Values.test.",
+			expected:    TemplateContext{"Values", "test"},
+		},
+		{
+			desc:        "Selects selector expression with $ correctly",
+			template:    `{{ $.Values.test }}`,
+			nodeContent: "$.Values.test",
+			expected:    TemplateContext{"$", "Values", "test"},
+		},
+		{
+			desc:        "Selects unfinished selector expression with $ correctly",
+			template:    `{{ $.Values.test. }}`,
+			nodeContent: "$.Values.test.",
+			expected:    TemplateContext{"$", "Values", "test"},
+		},
 		{
 			desc:        "Selects selector expression with variable correctly",
 			template:    `{{ $x.test }}`,
 			nodeContent: "$x.test",
 			expected:    TemplateContext{"$x", "test"},
 		},
+		{
+			desc:        "Selects unfinished selector expression with variable correctly",
+			template:    `{{ $x.test. }}`,
+			nodeContent: "$x.test.",
+			expected:    TemplateContext{"$x", "test"},
+		},
 	}
 	for _, tC := range testCases {
 		t.Run(tC.desc, func(t *testing.T) {

diff --git a/internal/lsp/symbol_table_test.go b/internal/lsp/symbol_table_test.go
@@ -217,9 +217,10 @@ func TestSymbolTableForValuesTestFile(t *testing.T) {
 
 func TestSymbolTableForValuesSingleTests(t *testing.T) {
 	type testCase struct {
-		template   string
-		path       []string
-		startPoint sitter.Point
+		template         string
+		path             []string
+		startPoint       sitter.Point
+		foundContextsLen int
 	}
 
 	testCases := []testCase{
@@ -236,6 +237,7 @@ func TestSymbolTableForValuesSingleTests(t *testing.T) {
 				Row:    2,
 				Column: 40,
 			},
+			foundContextsLen: 7,
 		},
 		{
 			template: `{{ $x := .Values }}{{ $x.test }}{{ .Values.test }}`,
@@ -244,6 +246,25 @@ func TestSymbolTableForValuesSingleTests(t *testing.T) {
 				Row:    0,
 				Column: 25,
 			},
+			foundContextsLen: 4,
+		},
+		{
+			template: `{{ $x.test }}`,
+			path:     []string{"$x", "test"},
+			startPoint: sitter.Point{
+				Row:    0,
+				Column: 6,
+			},
+			foundContextsLen: 2,
+		},
+		{
+			template: `{{ $x.test. }}`,
+			path:     []string{"$x", "test", ""},
+			startPoint: sitter.Point{
+				Row:    0,
+				Column: 10,
+			},
+			foundContextsLen: 3,
 		},
 		{
 			template: `{{ if (and .Values. ) }} {{ end }} `,
@@ -252,6 +273,7 @@ func TestSymbolTableForValuesSingleTests(t *testing.T) {
 				Row:    0,
 				Column: 12,
 			},
+			foundContextsLen: 2,
 		},
 		{
 			template: `{{ if (and .Values. ) }} {{ end }} `,
@@ -260,6 +282,7 @@ func TestSymbolTableForValuesSingleTests(t *testing.T) {
 				Row:    0,
 				Column: 18,
 			},
+			foundContextsLen: 2,
 		},
 	}
 
@@ -272,5 +295,6 @@ func TestSymbolTableForValuesSingleTests(t *testing.T) {
 			points = append(points, v.StartPoint)
 		}
 		assert.Contains(t, points, v.startPoint)
+		assert.Len(t, symbolTable.contexts, v.foundContextsLen)
 	}
 }
diff --git a/internal/lsp/visitor.go b/internal/lsp/visitor.go
@@ -43,6 +43,8 @@ func (v *Visitors) visitNodesRecursiveWithScopeShift(node *sitter.Node) {
 	case gotemplate.NodeTypeRangeAction:
 		rangeNode := node.ChildByFieldName("range")
 		if rangeNode == nil {
+			// for {{- range $type, $config := $root.Values.deployments }} the range node is in the
+			// range_variable_definition node an not in the range_action node
 			rangeNode = node.NamedChild(0).ChildByFieldName("range")
 			if rangeNode == nil {
 				logger.Error("Could not find range node")

diff --git a/testdata/example/templates/deployment.yaml b/testdata/example/templates/deployment.yaml
@@ -79,6 +79,6 @@ spec:
         name: my-app-{{ $type }}
       spec:
         replicas: {{ $config.hpa.minReplicas }}
-          {{ $.Values.ingress.hosts }}
+        test:  {{ $.Values.ingress.hosts }}
       ---
       {{- end }}