Add logic to process header objects of responses

wso2 · Nov 4, 2024 · b4270f2 · b4270f2
1 parent c849771
commit b4270f2
Showing 1 changed file with 26 additions and 10 deletions.
diff --git a/...rbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/definitions/OASParserUtil.java b/...rbon.apimgt.impl/src/main/java/org/wso2/carbon/apimgt/impl/definitions/OASParserUtil.java
@@ -654,6 +654,13 @@ private static void setRefOfApiResponse(ApiResponse response, SwaggerUpdateConte
                     addToReferenceObjectMap(ref, context);
                 }
             }
+
+            Map<String, Header> headers = response.getHeaders();
+            if (headers != null) {
+                for (Header header : headers.values()) {
+                    setRefOfApiResponseHeader(header, context);
+                }
+            }
         }
     }
 
@@ -683,9 +690,14 @@ private static void setRefOfApiResponseHeader(Header header, SwaggerUpdateContex
             if (content != null) {
                 extractReferenceFromContent(content, context);
             } else {
-                String ref = header.get$ref();
-                if (ref != null) {
-                    addToReferenceObjectMap(ref, context);
+                Schema schema = header.getSchema();
+                if (schema != null) {
+                    extractReferenceFromSchema(schema, context);
+                } else {
+                    String ref = header.get$ref();
+                    if (ref != null) {
+                        addToReferenceObjectMap(ref, context);
+                    }
                 }
             }
         }
@@ -781,15 +793,15 @@ private static void extractReferenceFromSchema(Schema schema, SwaggerUpdateConte
                         extractReferenceFromSchema(itemsSchema, context);
                     }
                 } else if (schema instanceof ObjectSchema) {
-                    references = addSchemaOfSchema(schema);
+                    references = addSchemaOfSchema(schema, context);
                 } else if (schema instanceof MapSchema) {
                     Schema additionalPropertiesSchema = (Schema) schema.getAdditionalProperties();
                     extractReferenceFromSchema(additionalPropertiesSchema, context);
                 } else if (schema instanceof ComposedSchema) {
                     if (((ComposedSchema) schema).getAllOf() != null) {
                         for (Schema sc : ((ComposedSchema) schema).getAllOf()) {
                             if (OBJECT_DATA_TYPE.equalsIgnoreCase(sc.getType())) {
-                                references.addAll(addSchemaOfSchema(sc));
+                                references.addAll(addSchemaOfSchema(sc, context));
                             } else {
                                 String schemaRef = sc.get$ref();
                                 if (schemaRef != null) {
@@ -803,7 +815,7 @@ private static void extractReferenceFromSchema(Schema schema, SwaggerUpdateConte
                     if (((ComposedSchema) schema).getAnyOf() != null) {
                         for (Schema sc : ((ComposedSchema) schema).getAnyOf()) {
                             if (OBJECT_DATA_TYPE.equalsIgnoreCase(sc.getType())) {
-                                references.addAll(addSchemaOfSchema(sc));
+                                references.addAll(addSchemaOfSchema(sc, context));
                             } else {
                                 String schemaRef = sc.get$ref();
                                 if (schemaRef != null) {
@@ -817,7 +829,7 @@ private static void extractReferenceFromSchema(Schema schema, SwaggerUpdateConte
                     if (((ComposedSchema) schema).getOneOf() != null) {
                         for (Schema sc : ((ComposedSchema) schema).getOneOf()) {
                             if (OBJECT_DATA_TYPE.equalsIgnoreCase(sc.getType())) {
-                                references.addAll(addSchemaOfSchema(sc));
+                                references.addAll(addSchemaOfSchema(sc, context));
                             } else {
                                 String schemaRef = sc.get$ref();
                                 if (schemaRef != null) {
@@ -866,13 +878,14 @@ private static void processSchemaProperties(Schema schema, SwaggerUpdateContext
         }
     }
 
-    private static List<String> addSchemaOfSchema(Schema schema) {
+    private static List<String> addSchemaOfSchema(Schema schema, SwaggerUpdateContext context) {
         List<String> references = new ArrayList<String>();
         ObjectSchema os = (ObjectSchema) schema;
         if (os.getProperties() != null) {
             for (String propertyName : os.getProperties().keySet()) {
-                if (os.getProperties().get(propertyName) instanceof ComposedSchema) {
-                    ComposedSchema cs = (ComposedSchema) os.getProperties().get(propertyName);
+                Schema propertySchema = os.getProperties().get(propertyName);
+                if (propertySchema instanceof ComposedSchema) {
+                    ComposedSchema cs = (ComposedSchema) propertySchema;
                     if (cs.getAllOf() != null) {
                         for (Schema sc : cs.getAllOf()) {
                             references.add(sc.get$ref());
@@ -888,6 +901,9 @@ private static List<String> addSchemaOfSchema(Schema schema) {
                     } else {
                         log.error("Unidentified schema. The schema is not available in the API definition.");
                     }
+                } else if (propertySchema instanceof ObjectSchema ||
+                        propertySchema instanceof ArraySchema) {
+                    extractReferenceFromSchema(propertySchema, context);
                 }
             }
         }