cleanup

internal/adapter/yamlls/diagnostics_integration_test.go

@@ -52,7 +52,7 @@ func (proc readWriteCloseMock) Close() error {
 }
 
 func readTestFiles(dir string, channel chan<- lsp.DidOpenTextDocumentParams, doneChan chan<- int) {
-	libRegEx, e := regexp.Compile(".*/templates/.*\\.yaml")
+	libRegEx, e := regexp.Compile(".*/templates/.*\\.ya?ml")
 	if e != nil {
 		log.Fatal(e)
 		return
@@ -81,20 +81,24 @@ func readTestFiles(dir string, channel chan<- lsp.DidOpenTextDocumentParams, don
 	doneChan <- count
 }
 
-func sendTestFilesToYamlls(documents *lsplocal.DocumentStore, yamllsConnector *Connector, doneInputChan <-chan int, doneChan chan<- int, channel <-chan lsp.DidOpenTextDocumentParams) {
+func sendTestFilesToYamlls(documents *lsplocal.DocumentStore, yamllsConnector *Connector,
+	doneReadingFilesChan <-chan int,
+	doneSendingFilesChan chan<- int,
+	filesChan <-chan lsp.DidOpenTextDocumentParams,
+) {
 	ownCount := 0
 	for {
 		select {
-		case d := <-channel:
+		case d := <-filesChan:
 			documents.DidOpen(d, util.DefaultConfig)
 			tree := lsplocal.ParseAst(nil, d.TextDocument.Text)
 			yamllsConnector.DocumentDidOpen(tree, d)
 			ownCount++
-		case count := <-doneInputChan:
+		case count := <-doneReadingFilesChan:
 			if count != ownCount {
 				log.Fatal("Count mismatch: ", count, " != ", ownCount)
 			}
-			doneChan <- count
+			doneSendingFilesChan <- count
 			return
 		}
 	}
@@ -111,6 +115,8 @@ func TestYamllsDiagnosticsIntegration(t *testing.T) {
 	config := util.DefaultConfig.YamllsConfiguration
 
 	yamllsSettings := util.DefaultYamllsSettings
+	// disabling yamlls schema store improves performance and
+	// removes all schema diagnostics that are not caused by the yaml trimming
 	yamllsSettings.Schemas = make(map[string]string)
 	yamllsSettings.YamllsSchemaStoreSettings = util.YamllsSchemaStoreSettings{
 		Enable: false,
@@ -122,14 +128,13 @@ func TestYamllsDiagnosticsIntegration(t *testing.T) {
 
 	didOpenChan := make(chan lsp.DidOpenTextDocumentParams)
 	go readTestFiles(TEST_DATA_DIR, didOpenChan, doneReadingFilesChan)
-	go sendTestFilesToYamlls(documents, yamllsConnector, doneReadingFilesChan, doneSendingFilesChan, didOpenChan)
+	go sendTestFilesToYamlls(documents,
+		yamllsConnector, doneReadingFilesChan, doneSendingFilesChan, didOpenChan)
 
-	sentCount := 0
-	receivedDiagnostics :=
-		make(map[uri.URI]lsp.PublishDiagnosticsParams)
-	diagnosticsCount := 0
+	sentCount, diagnosticsCount := 0, 0
+	receivedDiagnostics := make(map[uri.URI]lsp.PublishDiagnosticsParams)
 
-	afterCh := time.After(300 * time.Second)
+	afterCh := time.After(600 * time.Second)
 	for {
 		if sentCount != 0 && len(receivedDiagnostics) == sentCount {
 			fmt.Println("All files checked")
@@ -143,14 +148,13 @@ func TestYamllsDiagnosticsIntegration(t *testing.T) {
 				fmt.Printf("Got diagnostic in %s diagnostics: %v \n", d.URI.Filename(), d.Diagnostics)
 			}
 		case <-afterCh:
-			fmt.Println("Time's up! Checked ", len(receivedDiagnostics), " diagnostics.")
-			return
+			t.Fatal("Timed out waiting for diagnostics")
 		case count := <-doneSendingFilesChan:
 			sentCount = count
-			fmt.Println("Checked ", sentCount, " files")
 		}
 	}
 
 	fmt.Printf("Checked %d files, found %d diagnostics\n", sentCount, diagnosticsCount)
 	assert.Equal(t, diagnosticsCount, 22)
+	assert.Equal(t, 2368, sentCount, "Count of files in test data not equal to actual count")
 }

internal/adapter/yamlls/documentSync.go

@@ -28,7 +28,7 @@ func (yamllsConnector Connector) DocumentDidOpen(ast *sitter.Tree, params lsp.Di
 	if yamllsConnector.Conn == nil {
 		return
 	}
-	params.TextDocument.Text = trimTemplateForYamllsFromAst(ast, params.TextDocument.Text)
+	params.TextDocument.Text = lsplocal.TrimTemplate(ast, params.TextDocument.Text)
 
 	err := (*yamllsConnector.Conn).Notify(context.Background(), lsp.MethodTextDocumentDidOpen, params)
 	if err != nil {
@@ -40,7 +40,7 @@ func (yamllsConnector Connector) DocumentDidSave(doc *lsplocal.Document, params
 	if yamllsConnector.Conn == nil {
 		return
 	}
-	params.Text = trimTemplateForYamllsFromAst(doc.Ast, doc.Content)
+	params.Text = lsplocal.TrimTemplate(doc.Ast, doc.Content)
 
 	err := (*yamllsConnector.Conn).Notify(context.Background(), lsp.MethodTextDocumentDidSave, params)
 	if err != nil {
@@ -58,7 +58,7 @@ func (yamllsConnector Connector) DocumentDidChange(doc *lsplocal.Document, param
 	if yamllsConnector.Conn == nil {
 		return
 	}
-	trimmedText := trimTemplateForYamllsFromAst(doc.Ast, doc.Content)
+	trimmedText := lsplocal.TrimTemplate(doc.Ast, doc.Content)
 
 	logger.Debug("Sending DocumentDidChange previous", params)
 	for i, change := range params.ContentChanges {
@@ -89,7 +89,7 @@ func (yamllsConnector Connector) DocumentDidChangeFullSync(doc *lsplocal.Documen
 	}
 
 	logger.Println("Sending DocumentDidChange with full sync, current content:", doc.Content)
-	trimmedText := trimTemplateForYamllsFromAst(doc.Ast.Copy(), doc.Content)
+	trimmedText := lsplocal.TrimTemplate(doc.Ast.Copy(), doc.Content)
 
 	params.ContentChanges = []lsp.TextDocumentContentChangeEvent{
 		{

internal/lsp/yaml_ast.go

@@ -48,17 +48,15 @@ func TrimTemplate(gotemplateTree *sitter.Tree, content string) string {
 	ranges := getTextNodeRanges(gotemplateTree.RootNode())
 	result := make([]byte, len(content))
 	for i := range result {
-		if content[i] == '\n' {
-			result[i] = byte('\n')
-			continue
-		} else if content[i] == '\r' {
-			result[i] = byte('\r')
+		if content[i] == '\n' || content[i] == '\r' {
+			result[i] = content[i]
 			continue
 		}
 		result[i] = byte(' ')
 	}
 	for _, yamlRange := range ranges {
-		copy(result[yamlRange.StartByte:yamlRange.EndByte], content[yamlRange.StartByte:yamlRange.EndByte])
+		copy(result[yamlRange.StartByte:yamlRange.EndByte],
+			content[yamlRange.StartByte:yamlRange.EndByte])
 	}
 	return string(result)
 }