Restructure Code to Improve Code Quality

core/pkg/fixhandler/fixhandler.go

@@ -186,39 +186,34 @@ func (h *FixHandler) PrintExpectedChanges(resourcesToFix []ResourceFixInfo) {
 func (h *FixHandler) ApplyChanges(resourcesToFix []ResourceFixInfo) (int, []error) {
 	updatedFiles := make(map[string]bool)
 	errors := make([]error, 0)
-	// Map with key as filepath
-	filePathFixInfo := make(map[string]*fileFixInfo)
-	for _, resourceToFix := range resourcesToFix {
-		singleExpression := reduceYamlExpressions(&resourceToFix)
-		resourceFilePath := resourceToFix.FilePath
 
-		if _, pathExistsInMap := filePathFixInfo[resourceFilePath]; !pathExistsInMap {
-			filePathFixInfo[resourceFilePath] = &fileFixInfo{
-				contentToAdd:  make([]contentToAdd, 0),
-				linesToRemove: make([]linesToRemove, 0),
-			}
-		}
+	fileYamlExpressions := h.getFileYamlExpressions(resourcesToFix)
 
-		contentsToAdd, linesToRemove, err := h.getResourceFileFix(resourceFilePath, singleExpression, resourceToFix.DocumentIndex)
+	for filepath, yamlExpression := range fileYamlExpressions {
+		fileAsString, err := getFileString(filepath)
 
 		if err != nil {
-			errors = append(errors,
-				fmt.Errorf("failed to fix resource [Name: '%s', Kind: '%s'] in '%s': %w ",
-					resourceToFix.Resource.GetName(),
-					resourceToFix.Resource.GetKind(),
-					resourceToFix.FilePath,
-					err))
+			logger.L().Error(err.Error())
+			continue
+		}
+
+		fixedYamlString, err := h.ApplyFix(fileAsString, yamlExpression)
+
+		if err != nil {
+			errors = append(errors, fmt.Errorf("failed to fix file %s: %w ", filepath, err))
+			continue
 		} else {
-			h.addResourceFileFix(contentsToAdd, linesToRemove, filePathFixInfo[resourceFilePath])
-			updatedFiles[resourceToFix.FilePath] = true
+			updatedFiles[filepath] = true
+		}
+
+		err = writeFixesToFile(filepath, fixedYamlString)
+
+		if err != nil {
+			logger.L().Error(fmt.Sprintf("Cannot Apply fixes to file %s, %v", filepath, err.Error()))
+			errors = append(errors, err)
 		}
 	}
-	err := h.applyFixToFiles(filePathFixInfo)
 
-	if err != nil {
-		logger.L().Fatal(fmt.Sprintf("Cannot Apply fixes to files, %v", err.Error()))
-		errors = append(errors, err)
-	}
 	return len(updatedFiles), errors
 }
 
@@ -236,42 +231,40 @@ func (h *FixHandler) getFilePathAndIndex(filePathWithIndex string) (filePath str
 	}
 }
 
-func (h *FixHandler) getResourceFileFix(filePath string, yamlExpression string, documentIdx int) (*[]contentToAdd, *[]linesToRemove, error) {
-	originalYamlNode := (*constructDecodedYaml(filePath))[documentIdx]
-	fixedYamlNodes, err := constructFixedYamlNodes(filePath, yamlExpression)
+func (h *FixHandler) ApplyFix(yamlString, yamlExpression string) (fixedYamlString string, err error) {
+	yamlLines := strings.Split(yamlString, "\n")
+
+	originalRootNodes := constructDecodedYaml(yamlString)
+	fixedRootNodes, err := constructFixedYamlNodes(yamlString, yamlExpression)
+
 	if err != nil {
-		return nil, nil, err
+		return "", err
 	}
 
-	fixedYamlNode := (*fixedYamlNodes)[documentIdx]
+	contentsToAdd, linesToRemove := getFixInfo(originalRootNodes, fixedRootNodes)
 
-	originalList := constructDFSOrder(&originalYamlNode)
-	fixedList := constructDFSOrder(&fixedYamlNode)
+	fixedYamlLines := getFixedYamlLines(yamlLines, contentsToAdd, linesToRemove)
 
-	contentsToAdd, linesToRemove := getFixInfo(originalList, fixedList)
-
-	return contentsToAdd, linesToRemove, nil
+	fixedYamlString = getFixedYamlString(fixedYamlLines)
 
+	return fixedYamlString, nil
 }
 
-func (h *FixHandler) addResourceFileFix(contentToAdd *[]contentToAdd, linesToRemove *[]linesToRemove, fileFixInfo *fileFixInfo) {
-	for _, content := range *contentToAdd {
-		fileFixInfo.addContent(content)
-	}
+func (h *FixHandler) getFileYamlExpressions(resourcesToFix []ResourceFixInfo) map[string]string {
+	fileYamlExpressions := make(map[string]string, 0)
+	for _, resourceToFix := range resourcesToFix {
+		singleExpression := reduceYamlExpressions(&resourceToFix)
+		resourceFilePath := resourceToFix.FilePath
 
-	for _, lines := range *linesToRemove {
-		fileFixInfo.addLinesToRemove(lines)
-	}
-}
-
-func (h *FixHandler) applyFixToFiles(filePathFixInfo map[string]*fileFixInfo) error {
-	for filepath, fixInfo := range filePathFixInfo {
-		err := applyFixesToFile(filepath, &fixInfo.contentToAdd, &fixInfo.linesToRemove)
-		if err != nil {
-			return err
+		if _, pathExistsInMap := fileYamlExpressions[resourceFilePath]; !pathExistsInMap {
+			fileYamlExpressions[resourceFilePath] = singleExpression
+		} else {
+			fileYamlExpressions[resourceFilePath] = joinStrings(fileYamlExpressions[resourceFilePath], " | ", singleExpression)
 		}
+
 	}
-	return nil
+
+	return fileYamlExpressions
 }
 
 func (rfi *ResourceFixInfo) addYamlExpressionsFromResourceAssociatedControl(documentIndex int, ac *resourcesresults.ResourceAssociatedControl, skipUserValues bool) {
@@ -304,6 +297,30 @@ func reduceYamlExpressions(resource *ResourceFixInfo) string {
 	return strings.Join(expressions, " | ")
 }
 
+func joinStrings(inputStrings ...string) string {
+	return strings.Join(inputStrings, "")
+}
+
+func getFileString(filepath string) (string, error) {
+	bytes, err := ioutil.ReadFile(filepath)
+
+	if err != nil {
+		return "", fmt.Errorf("Error reading file %s", filepath)
+	}
+
+	return string(bytes), nil
+}
+
+func writeFixesToFile(filepath, content string) error {
+	err := ioutil.WriteFile(filepath, []byte(content), 0644)
+
+	if err != nil {
+		return fmt.Errorf("Error writing fixes to file: %w", err)
+	}
+
+	return nil
+}
+
 func fixPathToValidYamlExpression(fixPath, value string, documentIndexInYaml int) string {
 	isStringValue := true
 	if _, err := strconv.ParseBool(value); err == nil {

core/pkg/fixhandler/fixhandler_test.go

@@ -1,14 +1,10 @@
 package fixhandler
 
 import (
-	"fmt"
-	"io/ioutil"
 	"os"
 	"path/filepath"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-
 	logger "github.com/kubescape/go-logger"
 	metav1 "github.com/kubescape/kubescape/v2/core/meta/datastructures/v1"
 	reporthandlingv2 "github.com/kubescape/opa-utils/reporthandling/v2"
@@ -16,6 +12,12 @@ import (
 	"gopkg.in/op/go-logging.v1"
 )
 
+type indentationTestCase struct {
+	inputFile      string
+	yamlExpression string
+	expectedFile   string
+}
+
 func NewFixHandlerMock() (*FixHandler, error) {
 	backendLoggerLeveled := logging.AddModuleLevel(logging.NewLogBackend(logger.L().GetWriter(), "", 0))
 	backendLoggerLeveled.SetLevel(logging.ERROR, "")
@@ -33,151 +35,128 @@ func getTestdataPath() string {
 	return filepath.Join(currentDir, "testdata")
 }
 
-func testDirectoryApplyFixHelper(t *testing.T, yamlExpressions *[][]string, directoryPath string) {
+func getTestCases() []indentationTestCase {
+	indentationTestCases := []indentationTestCase{
+		// Insertion Scenarios
+		{
+			"insert_scenarios/original_yaml_scenario_1.yml",
+			"select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false",
+			"insert_scenarios/fixed_yaml_scenario_1.yml",
+		},
+		{
+			"insert_scenarios/original_yaml_scenario_2.yml",
+			"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"NET_RAW\"]",
+			"insert_scenarios/fixed_yaml_scenario_2.yml",
+		},
+		{
+			"insert_scenarios/original_yaml_scenario_3.yml",
+			"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]",
+			"insert_scenarios/fixed_yaml_scenario_3.yml",
+		},
+		{
+			"insert_scenarios/original_yaml_scenario_4.yml",
 
-	scenarioCount := len(*yamlExpressions)
+			`select(di==0).spec.template.spec.securityContext.allowPrivilegeEscalation |= false |
+			 select(di==0).spec.template.spec.containers[0].securityContext.capabilities.drop += ["NET_RAW"] |
+			 select(di==0).spec.template.spec.containers[0].securityContext.seccompProfile.type |= "RuntimeDefault" |
+			 select(di==0).spec.template.spec.containers[0].securityContext.allowPrivilegeEscalation |= false |
+			 select(di==0).spec.template.spec.containers[0].securityContext.readOnlyRootFilesystem |= true`,
 
-	for scenario := 1; scenario <= scenarioCount; scenario++ {
-		originalFile := fmt.Sprintf("original_yaml_scenario_%d.yml", scenario)
-		fixedFile := fmt.Sprintf("fixed_yaml_scenario_%d.yml", scenario)
+			"insert_scenarios/fixed_yaml_scenario_4.yml",
+		},
+		{
+			"insert_scenarios/original_yaml_scenario_5.yml",
+			"select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false",
+			"insert_scenarios/fixed_yaml_scenario_5.yml",
+		},
+		{
+			"insert_scenarios/original_yaml_scenario_6.yml",
+			"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]",
+			"insert_scenarios/fixed_yaml_scenario_6.yml",
+		},
+		{
+			"insert_scenarios/original_yaml_scenario_7.yml",
 
-		originalFilePath := filepath.Join(directoryPath, originalFile)
-		fixedFilePath := filepath.Join(directoryPath, fixedFile)
+			`select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false |
+			 select(di==1).spec.containers[0].securityContext.allowPrivilegeEscalation |= false`,
 
-		// create temp file
-		tempFile, err := ioutil.TempFile("", originalFile)
-		if err != nil {
-			panic(err)
-		}
-		defer os.Remove(tempFile.Name())
+			"insert_scenarios/fixed_yaml_scenario_7.yml",
+		},
 
-		// read original file
-		originalFileContent, err := ioutil.ReadFile(originalFilePath)
-		if err != nil {
-			panic(err)
+		// Removal Scenarios
+
+		{
+			"remove_scenarios/original_yaml_scenario_1.yml",
+			"del(select(di==0).spec.containers[0].securityContext)",
+			"remove_scenarios/fixed_yaml_scenario_1.yml",
+		},
+		{
+			"remove_scenarios/original_yaml_scenario_2.yml",
+			"del(select(di==0).spec.containers[1])",
+			"remove_scenarios/fixed_yaml_scenario_2.yml",
+		},
+		{
+			"remove_scenarios/original_yaml_scenario_3.yml",
+			"del(select(di==0).spec.containers[0].securityContext.capabilities.drop[1])",
+			"remove_scenarios/fixed_yaml_scenario_3.yml",
+		},
+		{
+			"remove_scenarios/original_yaml_scenario_4.yml",
+			`del(select(di==0).spec.containers[0].securityContext) | 
+			 del(select(di==1).spec.containers[1])`,
+			"remove_scenarios/fixed_yaml_scenario_4.yml",
+		},
+
+		// Replace Scenarios
+		{
+			"replace_scenarios/original_yaml_scenario_1.yml",
+			"select(di==0).spec.containers[0].securityContext.runAsRoot |= false",
+			"replace_scenarios/fixed_yaml_scenario_1.yml",
+		},
+		{
+			"replace_scenarios/original_yaml_scenario_2.yml",
+			`select(di==0).spec.containers[0].securityContext.capabilities.drop[0] |= "SYS_ADM" |
+			 select(di==0).spec.containers[0].securityContext.capabilities.add[0] |= "NET_RAW"`,
+			"replace_scenarios/fixed_yaml_scenario_2.yml",
+		},
+
+		// Hybrid Scenarios
+		{
+			"hybrid_scenarios/original_yaml_scenario_1.yml",
+			`del(select(di==0).spec.containers[0].securityContext) |
+			 select(di==0).spec.securityContext.runAsRoot |= false`,
+			"hybrid_scenarios/fixed_yaml_scenario_1.yml",
+		},
+	}
+
+	return indentationTestCases
+}
+
+func TestApplyFixKeepsIndentation(t *testing.T) {
+	testCases := getTestCases()
+
+	for _, tc := range testCases {
+		getTestDataPath := func(filename string) string {
+			currentDir, _ := os.Getwd()
+			currentFile := "testdata/" + filename
+			return filepath.Join(currentDir, currentFile)
 		}
 
-		// write original file contents to temp file
-		err = ioutil.WriteFile(tempFile.Name(), originalFileContent, 0644)
-		if err != nil {
-			panic(err)
-		}
+		input, _ := os.ReadFile(getTestDataPath(tc.inputFile))
+		want, _ := os.ReadFile(getTestDataPath(tc.expectedFile))
+		expression := tc.yamlExpression
 
-		// make changes to temp file
 		h, _ := NewFixHandlerMock()
 
-		filePathFixInfo := make(map[string]*fileFixInfo)
-		filePath := tempFile.Name()
-		filePathFixInfo[filePath] = &fileFixInfo{
-			contentToAdd:  make([]contentToAdd, 0),
-			linesToRemove: make([]linesToRemove, 0),
+		got, _ := h.ApplyFix(string(input), expression)
+
+		if got != string(want) {
+			t.Errorf("Fixed file does not match the expected.\nGot: <%s>\nWant:<%s>", got, want)
 		}
-		fixInfo := filePathFixInfo[filePath]
-
-		for idx, yamlExpression := range (*yamlExpressions)[scenario-1] {
-			contentToAdd, linesToRemove, err := h.getResourceFileFix(filePath, yamlExpression, idx)
-
-			if err == nil {
-				h.addResourceFileFix(contentToAdd, linesToRemove, fixInfo)
-			}
-
-		}
-
-		err = h.applyFixToFiles(filePathFixInfo)
-		assert.NoError(t, err)
-
-		// Check temp file contents
-		tempFileContent, err := ioutil.ReadFile(tempFile.Name())
-		if err != nil {
-			panic(err)
-		}
-
-		// Get fixed Yaml file content and check if it is equal to tempFileContent
-		fixedFileContent, err := ioutil.ReadFile(fixedFilePath)
-
-		errorMessage := fmt.Sprintf("Content of fixed %s doesn't match content of %s in %s", originalFile, fixedFile, directoryPath)
-
-		assert.Equal(t, string(fixedFileContent), string(tempFileContent), errorMessage)
 
 	}
 }
 
-func testDirectoryApplyFix(t *testing.T, directory string) {
-	directoryPath := filepath.Join(getTestdataPath(), directory)
-	var yamlExpressions [][]string
-
-	switch directory {
-	case "insert_scenarios":
-		yamlExpressions = [][]string{
-			{"select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false"},
-
-			{"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"NET_RAW\"]"},
-
-			{"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]"},
-
-			{`select(di==0).spec.template.spec.securityContext.allowPrivilegeEscalation |= false | 
-			 select(di==0).spec.template.spec.containers[0].securityContext.capabilities.drop += ["NET_RAW"] | 
-			 select(di==0).spec.template.spec.containers[0].securityContext.seccompProfile.type |= "RuntimeDefault" | 
-			 select(di==0).spec.template.spec.containers[0].securityContext.allowPrivilegeEscalation |= false | 
-			 select(di==0).spec.template.spec.containers[0].securityContext.readOnlyRootFilesystem |= true`},
-
-			{"select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false"},
-
-			{"select(di==0).spec.containers[0].securityContext.capabilities.drop += [\"SYS_ADM\"]"},
-
-			{
-				"select(di==0).spec.containers[0].securityContext.allowPrivilegeEscalation |= false",
-				"select(di==1).spec.containers[0].securityContext.allowPrivilegeEscalation |= false",
-			},
-		}
-
-	case "remove_scenarios":
-		yamlExpressions = [][]string{
-			{"del(select(di==0).spec.containers[0].securityContext)"},
-
-			{"del(select(di==0).spec.containers[1])"},
-
-			{"del(select(di==0).spec.containers[0].securityContext.capabilities.drop[1])"},
-
-			{
-				"del(select(di==0).spec.containers[0].securityContext)",
-				"del(select(di==1).spec.containers[1])",
-			},
-		}
-
-	case "replace_scenarios":
-		yamlExpressions = [][]string{
-			{"select(di==0).spec.containers[0].securityContext.runAsRoot |= false"},
-
-			{`select(di==0).spec.containers[0].securityContext.capabilities.drop[0] |= "SYS_ADM" | 
-			 select(di==0).spec.containers[0].securityContext.capabilities.add[0] |= "NET_RAW"`},
-		}
-
-	case "hybrid_scenarios":
-		yamlExpressions = [][]string{
-			{`del(select(di==0).spec.containers[0].securityContext) | 
-			 select(di==0).spec.securityContext.runAsRoot |= false`},
-		}
-	}
-
-	testDirectoryApplyFixHelper(t, &yamlExpressions, directoryPath)
-}
-
-func TestFixHandler_applyFixToFile(t *testing.T) {
-	// Tests for Insert scenarios
-	testDirectoryApplyFix(t, "insert_scenarios")
-
-	// Tests for Removal scenarios
-	testDirectoryApplyFix(t, "remove_scenarios")
-
-	// Tests for Replace scenarios
-	testDirectoryApplyFix(t, "replace_scenarios")
-
-	// Tests for Hybrid Scenarios
-	testDirectoryApplyFix(t, "hybrid_scenarios")
-
-}
-
 func Test_fixPathToValidYamlExpression(t *testing.T) {
 	type args struct {
 		fixPath             string

core/pkg/fixhandler/testdata/hybrid_scenarios/fixed_yaml_scenario_1.yml

@@ -16,4 +16,4 @@ spec:
   - name: nginx_container
     image: nginx
   securityContext:
-    runAsRoot: false
+    runAsRoot: false

core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_1.yml

@@ -11,4 +11,4 @@ spec:
   - name: nginx_container
     image: nginx
     securityContext:
-      allowPrivilegeEscalation: false
+      allowPrivilegeEscalation: false

core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_2.yml

@@ -12,4 +12,4 @@ spec:
     securityContext:
       capabilities:
         drop:
-          - NET_RAW
+          - NET_RAW

core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_3.yml

@@ -13,4 +13,4 @@ spec:
       capabilities:
         drop:
         - NET_RAW
-        - SYS_ADM
+        - SYS_ADM

core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_5.yml

@@ -15,4 +15,4 @@ spec:
 
   # Testing if comments are retained as intended
   securityContext:
-    runAsRoot: false
+    runAsRoot: false

core/pkg/fixhandler/testdata/insert_scenarios/fixed_yaml_scenario_6.yml

@@ -11,4 +11,4 @@ spec:
     image: nginx
     securityContext:
       capabilities:
-        drop: [NET_RAW, SYS_ADM]
+        drop: [NET_RAW, SYS_ADM]

core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_1.yml

@@ -9,4 +9,4 @@ metadata:
 spec:
   containers:
   - name: nginx_container
-    image: nginx
+    image: nginx

core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_2.yml

@@ -10,4 +10,3 @@ spec:
   containers:
   - name: nginx_container
     image: nginx
-

core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_3.yml

@@ -11,4 +11,4 @@ spec:
     image: nginx
     securityContext:
       capabilities:
-        drop: ["NET_RAW"]
+        drop: ["NET_RAW"]

core/pkg/fixhandler/testdata/remove_scenarios/fixed_yaml_scenario_4.yml

@@ -25,4 +25,3 @@ spec:
   containers:
   - name: nginx_container
     image: nginx
-

core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_1.yml

@@ -11,4 +11,4 @@ spec:
   - name: nginx_container
     image: nginx
     securityContext:
-      runAsRoot: false
+      runAsRoot: false

core/pkg/fixhandler/testdata/replace_scenarios/fixed_yaml_scenario_2.yml

@@ -15,4 +15,4 @@ spec:
       capabilities:
         drop:
         - "SYS_ADM"
-        add: ["NET_RAW"]
+        add: ["NET_RAW"]

core/pkg/fixhandler/yamlhandler.go

@@ -1,14 +1,11 @@
 package fixhandler
 
 import (
-	"bufio"
-	"bytes"
 	"container/list"
 	"errors"
 	"fmt"
 	"io"
-	"io/ioutil"
-	"os"
+	"strings"
 
 	logger "github.com/kubescape/go-logger"
 	"github.com/mikefarah/yq/v4/pkg/yqlib"
@@ -16,44 +13,37 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
-func constructDecodedYaml(filepath string) *[]yaml.Node {
-	file, err := ioutil.ReadFile(filepath)
-	if err != nil {
-		logger.L().Fatal("Cannot read file")
-	}
-	fileReader := bytes.NewReader(file)
+func constructDecodedYaml(yamlString string) *[]yaml.Node {
+	fileReader := strings.NewReader(yamlString)
 	dec := yaml.NewDecoder(fileReader)
 
 	nodes := make([]yaml.Node, 0)
 	for {
 		var node yaml.Node
-		err = dec.Decode(&node)
+		err := dec.Decode(&node)
+
 		nodes = append(nodes, node)
 
-		// break the loop in case of EOF
 		if errors.Is(err, io.EOF) {
 			break
 		}
 		if err != nil {
-			panic(err)
+			logger.L().Fatal("Cannot decode given document")
 		}
 	}
 
 	return &nodes
 }
 
-func constructFixedYamlNodes(filePath, yamlExpression string) (*[]yaml.Node, error) {
+func constructFixedYamlNodes(yamlString, yamlExpression string) (*[]yaml.Node, error) {
 	preferences := yqlib.ConfiguredYamlPreferences
 	preferences.EvaluateTogether = true
 	decoder := yqlib.NewYamlDecoder(preferences)
 
 	var allDocuments = list.New()
-	reader, err := constructNewReader(filePath)
-	if err != nil {
-		return nil, err
-	}
+	reader := strings.NewReader(yamlString)
 
-	fileDocuments, err := readDocuments(reader, filePath, 0, decoder)
+	fileDocuments, err := readDocuments(reader, decoder)
 	if err != nil {
 		return nil, err
 	}
@@ -103,7 +93,22 @@ func matchNodes(nodeOne, nodeTwo *yaml.Node) int {
 	}
 }
 
-func getFixInfo(originalList, fixedList *[]nodeInfo) (*[]contentToAdd, *[]linesToRemove) {
+func getFixInfo(originalRootNodes, fixedRootNodes *[]yaml.Node) (*[]contentToAdd, *[]linesToRemove) {
+	contentToAdd := make([]contentToAdd, 0)
+	linesToRemove := make([]linesToRemove, 0)
+
+	for idx, _ := range *fixedRootNodes {
+		originalList := constructDFSOrder(&(*originalRootNodes)[idx])
+		fixedList := constructDFSOrder(&(*fixedRootNodes)[idx])
+		nodeContentToAdd, nodeLinesToRemove := getFixInfoHelper(originalList, fixedList)
+		contentToAdd = append(contentToAdd, *nodeContentToAdd...)
+		linesToRemove = append(linesToRemove, *nodeLinesToRemove...)
+	}
+
+	return &contentToAdd, &linesToRemove
+}
+
+func getFixInfoHelper(originalList, fixedList *[]nodeInfo) (*[]contentToAdd, *[]linesToRemove) {
 
 	// While obtaining fixedYamlNode, comments and empty lines at the top are ignored.
 	// This causes a difference in Line numbers across the tree structure. In order to
@@ -232,71 +237,53 @@ func updateLinesToReplace(fixInfoMetadata *fixInfoMetadata) (int, int) {
 	return updatedOriginalTracker, updatedFixedTracker
 }
 
-func applyFixesToFile(filePath string, contentToAdd *[]contentToAdd, linesToRemove *[]linesToRemove) error {
-	// Read contents of the file line by line and store in a list
-	linesSlice, err := getLinesSlice(filePath)
-
-	if err != nil {
-		return err
-	}
-
-	// Determining last line required lineSlice. The placeholder for last line is replaced with the real last line
-	assignLastLine(contentToAdd, linesToRemove, &linesSlice)
-
-	// Clear the current content of file
-	if err := os.Truncate(filePath, 0); err != nil {
-		return err
-	}
-
-	file, err := os.OpenFile(filePath, os.O_RDWR, 0644)
-	if err != nil {
-		return err
-	}
-
-	defer func() error {
-		if err := file.Close(); err != nil {
-			return err
+func removeNewLinesAtTheEnd(yamlLines []string) []string {
+	for idx := 1; idx < len(yamlLines); idx++ {
+		if yamlLines[len(yamlLines)-idx] != "\n" {
+			yamlLines = yamlLines[:len(yamlLines)-idx+1]
+			break
 		}
-		return nil
-	}()
+	}
+	return yamlLines
+}
 
-	removeLines(linesToRemove, &linesSlice)
+func getFixedYamlLines(yamlLines []string, contentToAdd *[]contentToAdd, linesToRemove *[]linesToRemove) (fixedYamlLines []string) {
 
-	writer := bufio.NewWriter(file)
+	// Determining last line requires original yaml lines slice. The placeholder for last line is replaced with the real last line
+	assignLastLine(contentToAdd, linesToRemove, &yamlLines)
+
+	removeLines(linesToRemove, &yamlLines)
+
+	fixedYamlLines = make([]string, 0)
 	lineIdx, lineToAddIdx := 1, 0
 
 	// Ideally, new node is inserted at line before the next node in DFS order. But, when the previous line contains a
 	// comment or empty line, we need to insert new nodes before them.
-	adjustContentLines(contentToAdd, &linesSlice)
+	adjustContentLines(contentToAdd, &yamlLines)
 
 	for lineToAddIdx < len(*contentToAdd) {
 		for lineIdx <= (*contentToAdd)[lineToAddIdx].line {
 			// Check if the current line is not removed
-			if linesSlice[lineIdx-1] != "*" {
-				_, err := writer.WriteString(linesSlice[lineIdx-1] + "\n")
-				if err != nil {
-					return err
-				}
+			if yamlLines[lineIdx-1] != "*" {
+				fixedYamlLines = append(fixedYamlLines, yamlLines[lineIdx-1])
 			}
 			lineIdx += 1
 		}
 
 		content := (*contentToAdd)[lineToAddIdx].content
-		writer.WriteString(content)
+		fixedYamlLines = append(fixedYamlLines, content)
 
 		lineToAddIdx += 1
 	}
 
-	for lineIdx <= len(linesSlice) {
-		if linesSlice[lineIdx-1] != "*" {
-			_, err := writer.WriteString(linesSlice[lineIdx-1] + "\n")
-			if err != nil {
-				return err
-			}
+	for lineIdx <= len(yamlLines) {
+		if yamlLines[lineIdx-1] != "*" {
+			fixedYamlLines = append(fixedYamlLines, yamlLines[lineIdx-1])
 		}
 		lineIdx += 1
 	}
 
-	writer.Flush()
-	return nil
+	fixedYamlLines = removeNewLinesAtTheEnd(fixedYamlLines)
+
+	return fixedYamlLines
 }

core/pkg/fixhandler/yamlhelper.go

@@ -129,7 +129,7 @@ func constructContent(parentNode *yaml.Node, nodeList *[]nodeInfo, tracker int)
 
 	content = indentContent(content, indentationSpaces)
 
-	return content
+	return strings.TrimSuffix(content, "\n")
 }
 
 func indentContent(content string, indentationSpaces int) string {
@@ -306,12 +306,13 @@ func isEmptyLineOrComment(lineContent string) bool {
 	return false
 }
 
-func readDocuments(reader io.Reader, filename string, fileIndex int, decoder yqlib.Decoder) (*list.List, error) {
+func readDocuments(reader io.Reader, decoder yqlib.Decoder) (*list.List, error) {
 	err := decoder.Init(reader)
 	if err != nil {
 		return nil, err
 	}
 	inputList := list.New()
+
 	var currentIndex uint
 
 	for {
@@ -324,11 +325,10 @@ func readDocuments(reader io.Reader, filename string, fileIndex int, decoder yql
 			}
 			return inputList, nil
 		} else if errorReading != nil {
-			return nil, fmt.Errorf("bad file '%v': %w", filename, errorReading)
+			return nil, fmt.Errorf("Error Decoding YAML file")
 		}
+
 		candidateNode.Document = currentIndex
-		candidateNode.Filename = filename
-		candidateNode.FileIndex = fileIndex
 		candidateNode.EvaluateTogether = true
 
 		inputList.PushBack(candidateNode)
@@ -484,3 +484,7 @@ func updateTracker(nodeList *[]nodeInfo, tracker int) int {
 
 	return updatedTracker
 }
+
+func getFixedYamlString(yamlLines []string) (fixedYamlString string) {
+	return strings.Join(yamlLines, "\n")
+}