From 02cb395ec212de0b371b120cc10b61cf8519c4f8 Mon Sep 17 00:00:00 2001
From: Natasha Sarkar <natashasarkar@google.com>
Date: Thu, 2 Sep 2021 11:05:16 -0700
Subject: [PATCH] support multiple gvks in custom openapi schema

---
 api/krusty/openapicustomschema_test.go | 94 ++++++++++++++++++++++++++
 api/krusty/testdata/customschema.json  |  5 ++
 api/krusty/testdata/customschema.yaml  |  3 +
 kyaml/openapi/openapi.go               | 12 ++--
 4 files changed, 109 insertions(+), 5 deletions(-)

diff --git a/api/krusty/openapicustomschema_test.go b/api/krusty/openapicustomschema_test.go
index 94ffead18d..9e90e31d87 100644
--- a/api/krusty/openapicustomschema_test.go
+++ b/api/krusty/openapicustomschema_test.go
@@ -42,6 +42,26 @@ spec:
 `)
 }
 
+func writeOtherCustomResource(th kusttest_test.Harness, filepath string) {
+	th.WriteF(filepath, `
+apiVersion: v1alpha1
+kind: MyCRD
+metadata:
+  name: service
+spec:
+  template:
+    spec:
+      containers:
+      - name: server
+        image: server
+        command: example
+        ports:
+        - name: grpc
+          protocol: TCP
+          containerPort: 8080
+`)
+}
+
 func writeTestComponentWithCustomSchema(th kusttest_test.Harness) {
 	writeTestSchema(th, "comp/")
 	openapi.ResetOpenAPI()
@@ -74,6 +94,32 @@ patchesStrategicMerge:
           image: nginx
 `
 
+const customSchemaPatchMultipleGvks = `
+patchesStrategicMerge:
+- |-
+  apiVersion: example.com/v1alpha1
+  kind: MyCRD
+  metadata:
+    name: service
+  spec:
+    template:
+      spec:
+        containers:
+        - name: server
+          image: nginx
+- |-
+  apiVersion: v1alpha1
+  kind: MyCRD
+  metadata:
+    name: service
+  spec:
+    template:
+      spec:
+        containers:
+        - name: server
+          image: nginx
+`
+
 const patchedCustomResource = `
 apiVersion: example.com/v1alpha1
 kind: MyCRD
@@ -108,6 +154,54 @@ openapi:
 	th.AssertActualEqualsExpected(m, patchedCustomResource)
 }
 
+func TestCustomOpenApiFieldWithTwoGvks(t *testing.T) {
+	th := kusttest_test.MakeHarness(t)
+	th.WriteK(".", `
+resources:
+- mycrd.yaml
+- myothercrd.yaml
+openapi:
+  path: mycrd_schema.json
+`+customSchemaPatchMultipleGvks)
+	writeCustomResource(th, "mycrd.yaml")
+	writeOtherCustomResource(th, "myothercrd.yaml")
+	writeTestSchema(th, "./")
+	openapi.ResetOpenAPI()
+	m := th.Run(".", th.MakeDefaultOptions())
+	th.AssertActualEqualsExpected(m, `apiVersion: example.com/v1alpha1
+kind: MyCRD
+metadata:
+  name: service
+spec:
+  template:
+    spec:
+      containers:
+      - command: example
+        image: nginx
+        name: server
+        ports:
+        - containerPort: 8080
+          name: grpc
+          protocol: TCP
+---
+apiVersion: v1alpha1
+kind: MyCRD
+metadata:
+  name: service
+spec:
+  template:
+    spec:
+      containers:
+      - command: example
+        image: nginx
+        name: server
+        ports:
+        - containerPort: 8080
+          name: grpc
+          protocol: TCP
+`)
+}
+
 func TestCustomOpenApiFieldYaml(t *testing.T) {
 	th := kusttest_test.MakeHarness(t)
 	th.WriteK(".", `
diff --git a/api/krusty/testdata/customschema.json b/api/krusty/testdata/customschema.json
index eee65a97dd..04903e2255 100644
--- a/api/krusty/testdata/customschema.json
+++ b/api/krusty/testdata/customschema.json
@@ -34,6 +34,11 @@
           "group": "example.com",
           "kind": "MyCRD",
           "version": "v1alpha1"
+        },
+        {
+          "group": "",
+          "kind": "MyCRD",
+          "version": "v1alpha1"
         }
       ]
     },
diff --git a/api/krusty/testdata/customschema.yaml b/api/krusty/testdata/customschema.yaml
index 4e77065d4a..0f46cba927 100644
--- a/api/krusty/testdata/customschema.yaml
+++ b/api/krusty/testdata/customschema.yaml
@@ -22,6 +22,9 @@ definitions:
       - group: example.com
         kind: MyCRD
         version: v1alpha1
+      - group: ""
+        kind: MyCRD
+        version: v1alpha1
   io.k8s.api.core.v1.PodTemplateSpec:
     properties:
       metadata:
diff --git a/kyaml/openapi/openapi.go b/kyaml/openapi/openapi.go
index 5ad69adc70..9e6db90eb5 100644
--- a/kyaml/openapi/openapi.go
+++ b/kyaml/openapi/openapi.go
@@ -207,15 +207,17 @@ func AddDefinitions(definitions spec.Definitions) {
 		}
 		// cast the extension to a []map[string]string
 		exts, ok := gvk.([]interface{})
-		if !ok || len(exts) != 1 {
+		if !ok {
 			continue
 		}
 
-		typeMeta, ok := toTypeMeta(exts[0])
-		if !ok {
-			continue
+		for i := range exts {
+			typeMeta, ok := toTypeMeta(exts[i])
+			if !ok {
+				continue
+			}
+			globalSchema.schemaByResourceType[typeMeta] = &d
 		}
-		globalSchema.schemaByResourceType[typeMeta] = &d
 	}
 }