diff --git a/cli/cmd/policies_test.go b/cli/cmd/policies_test.go
index dc9ca4984a..7cce66eafc 100644
--- a/cli/cmd/policies_test.go
+++ b/cli/cmd/policies_test.go
@@ -70,7 +70,7 @@ func TestPoliciesFromKubeResources(t *testing.T) {
 		{
 			name: "valid input",
 			files: map[string]string{
-				"deployment.yaml": validDeploymentYAML,
+				"deployment.yml": validDeploymentYAML,
 			},
 			expectedOutput: []deployment{
 				{
@@ -83,21 +83,21 @@ func TestPoliciesFromKubeResources(t *testing.T) {
 		{
 			name: "missing annotation",
 			files: map[string]string{
-				"pod.yaml": podYAML,
+				"pod.yml": podYAML,
 			},
 		},
 		{
 			name: "invalid policy annotation",
 			files: map[string]string{
-				"deployment.yaml": invalidPolicyYAML,
+				"deployment.yml": invalidPolicyYAML,
 			},
 			expectedErr: "failed to parse policy test",
 		},
 		{
 			name: "multiple files",
 			files: map[string]string{
-				"deployment.yaml": validDeploymentYAML,
-				"pod.yaml":        anotherValidPodYAML,
+				"deployment.yml": validDeploymentYAML,
+				"pod.yml":        anotherValidPodYAML,
 			},
 			expectedOutput: []deployment{
 				{
diff --git a/cli/genpolicy/genpolicy_test.go b/cli/genpolicy/genpolicy_test.go
index 2d21888e0e..afb749a11e 100644
--- a/cli/genpolicy/genpolicy_test.go
+++ b/cli/genpolicy/genpolicy_test.go
@@ -55,7 +55,7 @@ func TestRunner(t *testing.T) {
 	expectedSettingsPath := "/settings.json"
 	settingsPathFile := filepath.Join(d, "settings_path")
 	cachePath := filepath.Join(d, "cache", "cache.json")
-	expectedYAMLPath := filepath.Join(d, "test.yaml")
+	expectedYAMLPath := filepath.Join(d, "test.yml")
 	yamlPathFile := filepath.Join(d, "yaml_path")
 
 	r, err := New(expectedRulesPath, expectedSettingsPath, cachePath, genpolicyBin)