Repository: FairwindsOps/polaris
Pull request #746 (2022-04-12 15:42:00)

Add fix command to mutate and update IaC 
This PR fixes #

## Checklist
* [ ] I have signed the CLA
* [ ] I have updated/added any relevant documentation
* [ ] Add the ability to add comments to the mutated file.

## Description
### What's the goal of this PR?
Add a fix command to fix IaC files that have failed checks and mutation defined. running something like `polaris fix --files-path test/mutate` should give you ability to modify all files under the `test/mutate` folder that failed and have mutation.

### What changes did you make?
- Added fix command
### What alternative solution should we consider, if any?




Patch:
diff --git a/checks/cpuLimitsMissing.yaml b/checks/cpuLimitsMissing.yaml
index 2e495524..bcb4a703 100644
--- a/checks/cpuLimitsMissing.yaml
+++ b/checks/cpuLimitsMissing.yaml
@@ -26,6 +26,12 @@ schema:
               not:
                 const: ''
 mutations:
+  - op: add
+    path: /resources/limits
+    value: {}
   - op: add
     path: /resources/limits/cpu
-    value: 100m
\ No newline at end of file
+    value: 100m
+comments:
+  - find: "cpu: 100m"
+    comment: "TODO: Set this to the amount of CPU you want to reserve for your workload"
\ No newline at end of file
diff --git a/checks/cpuRequestsMissing.yaml b/checks/cpuRequestsMissing.yaml
index cd1a56fb..8f2e8253 100644
--- a/checks/cpuRequestsMissing.yaml
+++ b/checks/cpuRequestsMissing.yaml
@@ -26,6 +26,12 @@ schema:
               not:
                 const: ''
 mutations:
+  - op: add
+    path: /resources/requests
+    value: {}
   - op: add
     path: /resources/requests/cpu
-    value: 100m
\ No newline at end of file
+    value: 100m
+comments:
+  - find: "cpu: 100m"
+    comment: "TODO: Set this to the amount of CPU you want to reserve for your workload"
\ No newline at end of file
diff --git a/checks/livenessProbeMissing.yaml b/checks/livenessProbeMissing.yaml
index 0b9ed043..f62319c3 100644
--- a/checks/livenessProbeMissing.yaml
+++ b/checks/livenessProbeMissing.yaml
@@ -22,4 +22,7 @@ schema:
 mutations:
   - op: add
     path: /livenessProbe
-    value: {"exec": { "command": [ "cat", "/tmp/healthy" ] }, "initialDelaySeconds": 5, "periodSeconds": 5 }
\ No newline at end of file
+    value: {"exec": { "command": [ "cat", "/tmp/healthy" ] }, "initialDelaySeconds": 5, "periodSeconds": 5 }
+comments:
+  - find: "livenessProbe:"
+    comment: "TODO: Change livenessProbe setting to reflect your health endpoints"
\ No newline at end of file
diff --git a/checks/memoryLimitsMissing.yaml b/checks/memoryLimitsMissing.yaml
index 431885a9..65288c13 100644
--- a/checks/memoryLimitsMissing.yaml
+++ b/checks/memoryLimitsMissing.yaml
@@ -31,4 +31,7 @@ mutations:
     value: {}
   - op: add
     path: /resources/limits/memory
-    value: "512Mi"
\ No newline at end of file
+    value: "512Mi"
+comments:
+  - find: "memory: 512Mi"
+    comment: "TODO: Set this to the amount of Memory you want to reserve for your workload"
\ No newline at end of file
diff --git a/checks/memoryRequestsMissing.yaml b/checks/memoryRequestsMissing.yaml
index 9051acfb..0e724ca2 100644
--- a/checks/memoryRequestsMissing.yaml
+++ b/checks/memoryRequestsMissing.yaml
@@ -31,4 +31,7 @@ mutations:
     value: {}
   - op: add
     path: /resources/requests/memory
-    value: "512Mi"
\ No newline at end of file
+    value: "512Mi"
+comments:
+  - find: "memory: 512Mi"
+    comment: "TODO: Set this to the amount of Memory you want to reserve for your workload"
\ No newline at end of file
diff --git a/checks/readinessProbeMissing.yaml b/checks/readinessProbeMissing.yaml
index 576370ef..4401e3b5 100644
--- a/checks/readinessProbeMissing.yaml
+++ b/checks/readinessProbeMissing.yaml
@@ -22,4 +22,7 @@ schema:
 mutations:
   - op: add
     path: /readinessProbe
-    value: {"exec": { "command": [ "cat", "/tmp/healthy" ] }, "initialDelaySeconds": 5, "periodSeconds": 5 }
\ No newline at end of file
+    value: {"exec": { "command": [ "cat", "/tmp/healthy" ] }, "initialDelaySeconds": 5, "periodSeconds": 5 }
+comments:
+  - find: "readinessProbe:"
+    comment: "TODO: Change livenessProbe setting to reflect your health endpoints"
\ No newline at end of file
diff --git a/cmd/polaris/fix.go b/cmd/polaris/fix.go
new file mode 100644
index 00000000..a5f7fe09
--- /dev/null
+++ b/cmd/polaris/fix.go
@@ -0,0 +1,146 @@
+// Copyright 2020 FairwindsOps Inc
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cmd
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+
+	"github.com/fairwindsops/polaris/pkg/kube"
+	"github.com/fairwindsops/polaris/pkg/mutation"
+	"github.com/fairwindsops/polaris/pkg/validator"
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/cobra"
+	yamlV3 "gopkg.in/yaml.v3"
+	"sigs.k8s.io/yaml"
+)
+
+var filesPath string
+
+func init() {
+	rootCmd.AddCommand(fixCommand)
+	fixCommand.PersistentFlags().StringVar(&filesPath, "files-path", "", "mutate and fix one or more YAML files in a specified folder")
+}
+
+var fixCommand = &cobra.Command{
+	Use:   "fix",
+	Short: "Fix Infrastructure as code files.",
+	Long:  `Fix Infrastructure as code files.`,
+	Run: func(cmd *cobra.Command, args []string) {
+		logrus.Debug("Setting up controller manager")
+
+		if filesPath == "" {
+			logrus.Error("Please specify a sub-command.")
+			err := cmd.Help()
+			panic(err)
+		}
+
+		baseDir := filesPath + "/"
+		yamlFiles, err := getYamlFiles(baseDir)
+		if err != nil {
+			panic(err)
+		}
+		var contentStr string
+		isFirstResource := true
+		for _, fullFilePath := range yamlFiles {
+
+			yamlFile, err := ioutil.ReadFile(fullFilePath)
+			if err != nil {
+				panic(err)
+			}
+
+			dec := yamlV3.NewDecoder(bytes.NewReader(yamlFile))
+
+			for {
+				data := map[string]interface{}{}
+				err := dec.Decode(&data)
+				// check it was parsed
+				if data == nil {
+					continue
+				}
+				// break the loop in case of EOF
+				if errors.Is(err, io.EOF) {
+					break
+				}
+				if err != nil {
+					panic(err)
+				}
+				yamlContent, err := yamlV3.Marshal(data)
+				if err != nil {
+					panic(err)
+				}
+				kubeResources := kube.CreateResourceProviderFromYaml(string(yamlContent))
+				results, err := validator.ApplyAllSchemaChecksToResourceProvider(&config, kubeResources)
+				if err != nil {
+					panic(err)
+				}
+				comments, allMutations := mutation.GetMutationsAndCommentsFromResults(results)
+				updatedYamlContent := string(yamlContent)
+				if len(allMutations) > 0 {
+					for _, resources := range kubeResources.Resources {
+						key := fmt.Sprintf("%s/%s/%s", resources[0].Kind, resources[0].Resource.GetName(), resources[0].Resource.GetNamespace())
+						mutations := allMutations[key]
+						mutated, err := mutation.ApplyAllSchemaMutations(&config, kubeResources, resources[0], mutations)
+						if err != nil {
+							panic(err)
+						}
+						mutatedYamlContent, err := yaml.JSONToYAML(mutated.OriginalObjectJSON)
+						if err != nil {
+							panic(err)
+						}
+						updatedYamlContent = mutation.UpdateMutatedContentWithComments(string(mutatedYamlContent), comments)
+					}
+				}
+				if isFirstResource {
+					contentStr = updatedYamlContent
+					isFirstResource = false
+				} else {
+					contentStr += "\n"
+					contentStr += "---"
+					contentStr += "\n"
+					contentStr += updatedYamlContent
+				}
+			}
+
+			if contentStr != "" {
+				err = ioutil.WriteFile(fullFilePath, []byte(contentStr), 0644)
+				if err != nil {
+					logrus.Errorf("Error writing output to file: %v", err)
+					os.Exit(1)
+				}
+			}
+		}
+
+	},
+}
+
+func getYamlFiles(rootpath string) ([]string, error) {
+	var list []string
+	err := filepath.Walk(rootpath, func(path string, info os.FileInfo, err error) error {
+		if info.IsDir() {
+			return nil
+		}
+		if filepath.Ext(path) == ".yaml" {
+			list = append(list, path)
+		}
+		return nil
+	})
+	return list, err
+}
diff --git a/examples/config.yaml b/examples/config.yaml
index a06a2dac..dadec568 100644
--- a/examples/config.yaml
+++ b/examples/config.yaml
@@ -27,8 +27,7 @@ checks:
   hostNetworkSet: danger
   hostPortSet: warning
   tlsSettingsMissing: warning
-
-
+ 
 exemptions:
   - namespace: kube-system
     controllerNames:
diff --git a/go.mod b/go.mod
index 3378f664..78c2d8af 100644
--- a/go.mod
+++ b/go.mod
@@ -19,6 +19,7 @@ require (
 	github.com/stretchr/testify v1.7.0
 	github.com/thoas/go-funk v0.9.1
 	golang.org/x/text v0.3.7 // indirect
+	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b
 	k8s.io/api v0.22.2
 	k8s.io/apimachinery v0.22.2
diff --git a/pkg/config/schema.go b/pkg/config/schema.go
index 678fc01f..692ced0a 100644
--- a/pkg/config/schema.go
+++ b/pkg/config/schema.go
@@ -49,6 +49,12 @@ var HandledTargets = []TargetKind{
 	TargetPod,
 }
 
+// MutationComment is the comments added to a mutated file
+type MutationComment struct {
+	Find    string `yaml:"find" json:"find"`
+	Comment string `yaml:"comment" json:"comment"`
+}
+
 // SchemaCheck is a Polaris check that runs using JSON Schema
 type SchemaCheck struct {
 	ID                      string                            `yaml:"id" json:"id"`
@@ -66,6 +72,7 @@ type SchemaCheck struct {
 	AdditionalSchemaStrings map[string]string                 `yaml:"additionalSchemaStrings" json:"additionalSchemaStrings"`
 	AdditionalValidators    map[string]jsonschema.RootSchema  `yaml:"-" json:"-"`
 	Mutations               []map[string]interface{}          `yaml:"mutations" json:"mutations"`
+	Comments                []MutationComment                 `yaml:"comments" json:"comments"`
 }
 
 type resourceMinimum string
diff --git a/pkg/kube/resource.go b/pkg/kube/resource.go
index 5c87f286..87d07ba3 100644
--- a/pkg/kube/resource.go
+++ b/pkg/kube/resource.go
@@ -40,6 +40,10 @@ type GenericResource struct {
 
 // NewGenericResourceFromUnstructured creates a workload from an unstructured.Unstructured
 func NewGenericResourceFromUnstructured(unst unstructured.Unstructured, podSpecMap interface{}) (GenericResource, error) {
+	if unst.GetCreationTimestamp().Time.IsZero() {
+		unstructured.RemoveNestedField(unst.Object, "metadata", "creationTimestamp")
+		unstructured.RemoveNestedField(unst.Object, "status")
+	}
 	workload := GenericResource{
 		Kind:     unst.GetKind(),
 		Resource: unst,
diff --git a/pkg/kube/resources.go b/pkg/kube/resources.go
index 9c4ecca3..51a92224 100644
--- a/pkg/kube/resources.go
+++ b/pkg/kube/resources.go
@@ -225,6 +225,13 @@ func CreateResourceProviderFromPath(directory string) (*ResourceProvider, error)
 	return &resources, nil
 }
 
+// CreateResourceProviderFromYaml returns a new ResourceProvider using the yaml
+func CreateResourceProviderFromYaml(yamlContent string) *ResourceProvider {
+	resources := newResourceProvider("unknown", "Content", "unknown")
+	resources.addResourcesFromYaml(string(yamlContent))
+	return &resources
+}
+
 // CreateResourceProviderFromCluster creates a new ResourceProvider using live data from a cluster
 func CreateResourceProviderFromCluster(ctx context.Context, c conf.Configuration) (*ResourceProvider, error) {
 	kubeConf, configError := config.GetConfigWithContext(c.KubeContext)
diff --git a/pkg/mutation/mutate.go b/pkg/mutation/mutate.go
index 7668ef71..df07b98e 100644
--- a/pkg/mutation/mutate.go
+++ b/pkg/mutation/mutate.go
@@ -1,13 +1,16 @@
 package mutation
 
 import (
+	"bufio"
 	"encoding/json"
 	"fmt"
+	"strings"
 
 	jsonpatch "github.com/evanphx/json-patch/v5"
 	"github.com/fairwindsops/polaris/pkg/config"
 	"github.com/fairwindsops/polaris/pkg/kube"
 	"github.com/fairwindsops/polaris/pkg/validator"
+	"github.com/thoas/go-funk"
 )
 
 // ApplyAllSchemaMutations applies available mutation to a single resource
@@ -35,9 +38,10 @@ func ApplyAllSchemaMutations(conf *config.Configuration, resourceProvider *kube.
 	return mutated, nil
 }
 
-// GetMutationsFromResults returns all mutations from results
-func GetMutationsFromResults(conf *config.Configuration, results []validator.Result) map[string][]map[string]interface{} {
+// GetMutationsAndCommentsFromResults returns all mutations from results
+func GetMutationsAndCommentsFromResults(results []validator.Result) ([]config.MutationComment, map[string][]map[string]interface{}) {
 	allMutationsFromResults := make(map[string][]map[string]interface{})
+	comments := []config.MutationComment{}
 	for _, result := range results {
 		key := fmt.Sprintf("%s/%s/%s", result.Kind, result.Name, result.Namespace)
 
@@ -49,6 +53,9 @@ func GetMutationsFromResults(conf *config.Configuration, results []validator.Res
 				}
 				allMutationsFromResults[key] = append(mutations, resultMessage.Mutations...)
 			}
+			if len(resultMessage.Comments) > 0 {
+				comments = append(comments, resultMessage.Comments...)
+			}
 		}
 
 		for _, resultMessage := range result.PodResult.Results {
@@ -59,6 +66,9 @@ func GetMutationsFromResults(conf *config.Configuration, results []validator.Res
 				}
 				allMutationsFromResults[key] = append(mutations, resultMessage.Mutations...)
 			}
+			if len(resultMessage.Comments) > 0 {
+				comments = append(comments, resultMessage.Comments...)
+			}
 		}
 
 		for _, containerResult := range result.PodResult.ContainerResults {
@@ -70,9 +80,36 @@ func GetMutationsFromResults(conf *config.Configuration, results []validator.Res
 					}
 					allMutationsFromResults[key] = append(mutations, resultMessage.Mutations...)
 				}
+				if len(resultMessage.Comments) > 0 {
+					comments = append(comments, resultMessage.Comments...)
+				}
 			}
 		}
 
 	}
-	return allMutationsFromResults
+	return comments, allMutationsFromResults
+}
+
+// UpdateMutatedContentWithComments Updates mutated object with comments
+func UpdateMutatedContentWithComments(yamlContent string, comments []config.MutationComment) string {
+	var lines []string
+	scanner := bufio.NewScanner(strings.NewReader(yamlContent))
+	scanner.Split(bufio.ScanLines)
+	for scanner.Scan() {
+		line := scanner.Text()
+		lines = append(lines, line)
+	}
+	commentMap := funk.Map(comments, func(c config.MutationComment) (string, string) {
+		return c.Find, c.Comment
+	}).(map[string]string)
+	fileContent := ""
+	for _, line := range lines {
+		comment, ok := commentMap[strings.TrimSpace(line)]
+		if ok {
+			line += (" #" + comment)
+		}
+		fileContent += line
+		fileContent += "\n"
+	}
+	return fileContent
 }
diff --git a/pkg/validator/output.go b/pkg/validator/output.go
index 94b476c8..08257fd9 100644
--- a/pkg/validator/output.go
+++ b/pkg/validator/output.go
@@ -79,6 +79,7 @@ type ResultMessage struct {
 	Severity  config.Severity
 	Category  string
 	Mutations []map[string]interface{}
+	Comments  []config.MutationComment
 }
 
 // ResultSet contiains the results for a set of checks
diff --git a/pkg/validator/schema.go b/pkg/validator/schema.go
index 16ca5f33..c37d1060 100644
--- a/pkg/validator/schema.go
+++ b/pkg/validator/schema.go
@@ -319,6 +319,7 @@ func applySchemaCheck(conf *config.Configuration, checkID string, test schemaTes
 				return mutationCopy
 			}).([]map[string]interface{})
 			result.Mutations = mutations
+			result.Comments = check.Comments
 		}
 	}
 	return &result, nil


Test patch:
diff --git a/test/checks/cpuLimitsMissing/mutated.yaml b/test/checks/cpuLimitsMissing/mutated.yaml
new file mode 100644
index 00000000..b56c750f
--- /dev/null
+++ b/test/checks/cpuLimitsMissing/mutated.yaml
@@ -0,0 +1,14 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    app.kubernetes.io/name: nginx
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    name: nginx
+    resources:
+      limits:
+        cpu: 100m #TODO: Set this to the amount of CPU you want to reserve for your workload
+        memory: 128Mi
diff --git a/test/checks/cpuRequestsMissing/mutated.yaml b/test/checks/cpuRequestsMissing/mutated.yaml
new file mode 100644
index 00000000..386cfc65
--- /dev/null
+++ b/test/checks/cpuRequestsMissing/mutated.yaml
@@ -0,0 +1,14 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    app.kubernetes.io/name: nginx
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    name: nginx
+    resources:
+      requests:
+        cpu: 100m #TODO: Set this to the amount of CPU you want to reserve for your workload
+        memory: 128Mi
diff --git a/test/checks/deploymentMissingReplicas/mutated.yaml b/test/checks/deploymentMissingReplicas/mutated.yaml
new file mode 100644
index 00000000..d11eca20
--- /dev/null
+++ b/test/checks/deploymentMissingReplicas/mutated.yaml
@@ -0,0 +1,11 @@
+apiVersion: v1
+kind: Deployment
+metadata:
+  labels:
+    env: test
+  name: nginx
+spec:
+  replicas: 3
+  template:
+    spec:
+      containers: []
diff --git a/test/checks/deploymentMissingReplicas/success.yaml b/test/checks/deploymentMissingReplicas/success.yaml
index 3c5d5269..cd005860 100644
--- a/test/checks/deploymentMissingReplicas/success.yaml
+++ b/test/checks/deploymentMissingReplicas/success.yaml
@@ -8,5 +8,4 @@ spec:
   replicas: 3
   template:
     spec:
-      containers: []
-
+      containers: []
\ No newline at end of file
diff --git a/test/checks/hostIPCSet/mutated.yaml b/test/checks/hostIPCSet/mutated.yaml
new file mode 100644
index 00000000..1b0b62e6
--- /dev/null
+++ b/test/checks/hostIPCSet/mutated.yaml
@@ -0,0 +1,11 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    app.kubernetes.io/name: nginx
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    name: nginx
+    resources: {}
diff --git a/test/checks/hostIPCSet/success.yaml b/test/checks/hostIPCSet/success.yaml
index 5af3ef74..8cac6060 100644
--- a/test/checks/hostIPCSet/success.yaml
+++ b/test/checks/hostIPCSet/success.yaml
@@ -8,5 +8,4 @@ spec:
   hostIPC: false
   containers:
   - name: nginx 
-    image: nginx
-        
\ No newline at end of file
+    image: nginx
\ No newline at end of file
diff --git a/test/checks/hostNetworkSet/mutated.yaml b/test/checks/hostNetworkSet/mutated.yaml
new file mode 100644
index 00000000..1b0b62e6
--- /dev/null
+++ b/test/checks/hostNetworkSet/mutated.yaml
@@ -0,0 +1,11 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    app.kubernetes.io/name: nginx
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    name: nginx
+    resources: {}
diff --git a/test/checks/hostPIDSet/mutated.yaml b/test/checks/hostPIDSet/mutated.yaml
new file mode 100644
index 00000000..1b0b62e6
--- /dev/null
+++ b/test/checks/hostPIDSet/mutated.yaml
@@ -0,0 +1,11 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    app.kubernetes.io/name: nginx
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    name: nginx
+    resources: {}
diff --git a/test/checks/livenessProbeMissing/mutated.yaml b/test/checks/livenessProbeMissing/mutated.yaml
new file mode 100644
index 00000000..0c8b43f0
--- /dev/null
+++ b/test/checks/livenessProbeMissing/mutated.yaml
@@ -0,0 +1,22 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    test: liveness
+  name: liveness-exec
+spec:
+  containers:
+  - args:
+    - /bin/sh
+    - -c
+    - touch /tmp/healthy; sleep 30; rm -rf /tmp/healthy; sleep 600
+    image: k8s.gcr.io/busybox
+    livenessProbe: #TODO: Change livenessProbe setting to reflect your health endpoints
+      exec:
+        command:
+        - cat
+        - /tmp/healthy
+      initialDelaySeconds: 5
+      periodSeconds: 5
+    name: liveness
+    resources: {}
diff --git a/test/checks/memoryLimitsMissing/mutated.yaml b/test/checks/memoryLimitsMissing/mutated.yaml
new file mode 100644
index 00000000..fbba0769
--- /dev/null
+++ b/test/checks/memoryLimitsMissing/mutated.yaml
@@ -0,0 +1,13 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    app.kubernetes.io/name: nginx
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    name: nginx
+    resources:
+      limits:
+        memory: 512Mi #TODO: Set this to the amount of Memory you want to reserve for your workload
diff --git a/test/checks/memoryRequestsMissing/mutated.yaml b/test/checks/memoryRequestsMissing/mutated.yaml
new file mode 100644
index 00000000..432dc1ab
--- /dev/null
+++ b/test/checks/memoryRequestsMissing/mutated.yaml
@@ -0,0 +1,13 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    app.kubernetes.io/name: nginx
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    name: nginx
+    resources:
+      requests:
+        memory: 512Mi #TODO: Set this to the amount of Memory you want to reserve for your workload
diff --git a/test/checks/pullPolicyNotAlways/mutated.cron-job.yaml b/test/checks/pullPolicyNotAlways/mutated.cron-job.yaml
new file mode 100644
index 00000000..67443f9b
--- /dev/null
+++ b/test/checks/pullPolicyNotAlways/mutated.cron-job.yaml
@@ -0,0 +1,19 @@
+apiVersion: batch/v1
+kind: CronJob
+metadata:
+  name: hello
+spec:
+  jobTemplate:
+    spec:
+      template:
+        spec:
+          containers:
+          - command:
+            - /bin/sh
+            - -c
+            - date; echo Hello from the Kubernetes cluster
+            image: busybox
+            imagePullPolicy: Always
+            name: hello
+          restartPolicy: OnFailure
+  schedule: '* * * * *'
diff --git a/test/checks/pullPolicyNotAlways/mutated.deployment.yaml b/test/checks/pullPolicyNotAlways/mutated.deployment.yaml
new file mode 100644
index 00000000..2855ec1d
--- /dev/null
+++ b/test/checks/pullPolicyNotAlways/mutated.deployment.yaml
@@ -0,0 +1,22 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  labels:
+    app: nginx
+  name: nginx-deployment
+spec:
+  replicas: 3
+  selector:
+    matchLabels:
+      app: nginx
+  template:
+    metadata:
+      labels:
+        app: nginx
+    spec:
+      containers:
+      - image: nginx:1.14.2
+        imagePullPolicy: Always
+        name: nginx
+        ports:
+        - containerPort: 80
diff --git a/test/checks/pullPolicyNotAlways/mutated.yaml b/test/checks/pullPolicyNotAlways/mutated.yaml
new file mode 100644
index 00000000..80d48d19
--- /dev/null
+++ b/test/checks/pullPolicyNotAlways/mutated.yaml
@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    env: test
+  name: nginx
+spec:
+  containers:
+  - image: nginx
+    imagePullPolicy: Always
+    name: nginx
+    resources: {}
diff --git a/test/checks/pullPolicyNotAlways/success.yaml b/test/checks/pullPolicyNotAlways/success.yaml
index 32a764d4..782149cc 100644
--- a/test/checks/pullPolicyNotAlways/success.yaml
+++ b/test/checks/pullPolicyNotAlways/success.yaml
@@ -8,5 +8,4 @@ spec:
   containers:
   - name: nginx
     image: nginx
-    imagePullPolicy: Always
-    
\ No newline at end of file
+    imagePullPolicy: Always
\ No newline at end of file
diff --git a/test/checks/readinessProbeMissing/mutated.yaml b/test/checks/readinessProbeMissing/mutated.yaml
new file mode 100644
index 00000000..97d3cbec
--- /dev/null
+++ b/test/checks/readinessProbeMissing/mutated.yaml
@@ -0,0 +1,22 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  labels:
+    test: readiness
+  name: readiness-exec
+spec:
+  containers:
+  - args:
+    - /bin/sh
+    - -c
+    - touch /tmp/healthy; sleep 30; rm -rf /tmp/healthy; sleep 600
+    image: k8s.gcr.io/busybox
+    name: readiness
+    readinessProbe: #TODO: Change livenessProbe setting to reflect your health endpoints
+      exec:
+        command:
+        - cat
+        - /tmp/healthy
+      initialDelaySeconds: 5
+      periodSeconds: 5
+    resources: {}
diff --git a/test/checks/runAsRootAllowed/mutated.yaml b/test/checks/runAsRootAllowed/mutated.yaml
new file mode 100644
index 00000000..d660e41b
--- /dev/null
+++ b/test/checks/runAsRootAllowed/mutated.yaml
@@ -0,0 +1,12 @@
+apiVersion: v1
+kind: Deployment
+metadata:
+  name: nginx
+spec:
+  template:
+    spec:
+      containers:
+      - image: nginx
+        name: nginx
+        securityContext:
+          runAsNonRoot: true
diff --git a/test/checks/runAsRootAllowed/success.yaml b/test/checks/runAsRootAllowed/success.yaml
index f163fb5b..0b45eaea 100644
--- a/test/checks/runAsRootAllowed/success.yaml
+++ b/test/checks/runAsRootAllowed/success.yaml
@@ -9,4 +9,4 @@ spec:
       - name: nginx
         image: nginx
         securityContext:
-          runAsNonRoot: true
+          runAsNonRoot: true
\ No newline at end of file
diff --git a/test/mutation_test.go b/test/mutation_test.go
index ba8d39da..a2212981 100644
--- a/test/mutation_test.go
+++ b/test/mutation_test.go
@@ -9,6 +9,7 @@ import (
 	"github.com/fairwindsops/polaris/pkg/mutation"
 	"github.com/fairwindsops/polaris/pkg/validator"
 	"github.com/stretchr/testify/assert"
+	"sigs.k8s.io/yaml"
 )
 
 var configYaml = `
@@ -36,26 +37,26 @@ func TestMutations(t *testing.T) {
 	for _, mutationStr := range mutations {
 		for _, tc := range failureTestCasesMap[mutationStr] {
 			newConfig := c
-			key := fmt.Sprintf("%s/%s", tc.check, strings.ReplaceAll(tc.filename, "failure", "success"))
-			successResources, ok := successResourceMap[key]
+			key := fmt.Sprintf("%s/%s", tc.check, strings.ReplaceAll(tc.filename, "failure", "mutated"))
+			mutatedYamlContent, ok := mutatedYamlContentMap[key]
 			assert.True(t, ok)
 			assert.Len(t, tc.resources.Resources, 1)
-			assert.Len(t, successResources.Resources, 1)
 			newConfig.Mutations = []string{mutationStr}
 			results, err := validator.ApplyAllSchemaChecksToResourceProvider(&newConfig, tc.resources)
 			assert.NoError(t, err)
 			assert.Len(t, results, 1)
-			allMutations := mutation.GetMutationsFromResults(&c, results)
+			comments, allMutations := mutation.GetMutationsAndCommentsFromResults(results)
 			assert.Len(t, allMutations, 1)
-			for kind, resources := range tc.resources.Resources {
+			for _, resources := range tc.resources.Resources {
 				assert.Len(t, resources, 1)
 				key := fmt.Sprintf("%s/%s/%s", resources[0].Kind, resources[0].Resource.GetName(), resources[0].Resource.GetNamespace())
 				mutations := allMutations[key]
 				mutated, err := mutation.ApplyAllSchemaMutations(&c, tc.resources, resources[0], mutations)
 				assert.NoError(t, err)
-				expected := successResources.Resources[kind][0]
-				// use fmt.sprint incase type casting issues for int and float.
-				assert.EqualValues(t, fmt.Sprint(expected.Resource.Object), fmt.Sprint(mutated.Resource.Object))
+				yamlContent, err := yaml.JSONToYAML(mutated.OriginalObjectJSON)
+				assert.NoError(t, err)
+				contentStr := mutation.UpdateMutatedContentWithComments(string(yamlContent), comments)
+				assert.EqualValues(t, mutatedYamlContent, contentStr)
 			}
 		}
 	}
diff --git a/test/schema_test.go b/test/schema_test.go
index ed9ba000..2aca7cae 100644
--- a/test/schema_test.go
+++ b/test/schema_test.go
@@ -17,6 +17,7 @@ package test
 import (
 	"fmt"
 	"io/ioutil"
+	"os"
 	"path/filepath"
 	"runtime"
 	"strings"
@@ -38,7 +39,7 @@ type testCase struct {
 	failure   bool
 }
 
-var successResourceMap = map[string]*kube.ResourceProvider{}
+var mutatedYamlContentMap = map[string]string{}
 var failureTestCasesMap = map[string][]testCase{}
 
 func init() {
@@ -60,24 +61,31 @@ func init() {
 			if err != nil {
 				panic(err)
 			}
-			testcase := testCase{
-				filename:  tc.Name(),
-				check:     check,
-				resources: resources,
-				failure:   strings.Contains(tc.Name(), "failure"),
-			}
-			testCases = append(testCases, testcase)
 
-			if strings.Contains(tc.Name(), "success") {
+			if strings.Contains(tc.Name(), "mutated") {
+				yamlContent, err := os.ReadFile(checkDir + "/" + tc.Name())
+				if err != nil {
+					panic(err)
+				}
 				key := fmt.Sprintf("%s/%s", check, tc.Name())
-				successResourceMap[key] = resources
+				mutatedYamlContentMap[key] = string(yamlContent)
 			} else {
-				testCases, ok := failureTestCasesMap[check]
-				if !ok {
-					testCases = []testCase{}
+				testcase := testCase{
+					filename:  tc.Name(),
+					check:     check,
+					resources: resources,
+					failure:   strings.Contains(tc.Name(), "failure"),
 				}
 				testCases = append(testCases, testcase)
-				failureTestCasesMap[check] = testCases
+
+				if strings.Contains(tc.Name(), "mutated") {
+					testCases, ok := failureTestCasesMap[check]
+					if !ok {
+						testCases = []testCase{}
+					}
+					testCases = append(testCases, testcase)
+					failureTestCasesMap[check] = testCases
+				}
 			}
 		}
 	}