From dd52ac840709123294b759feb143baa75c39b228 Mon Sep 17 00:00:00 2001
From: lnash94 <sumudunissanke@gmail.com>
Date: Wed, 25 Sep 2024 22:43:16 +0530
Subject: [PATCH] Separate to different functions

---
 .../core/generators/common/OASModifier.java   | 121 ++++++++++--------
 1 file changed, 70 insertions(+), 51 deletions(-)

diff --git a/openapi-core/src/main/java/io/ballerina/openapi/core/generators/common/OASModifier.java b/openapi-core/src/main/java/io/ballerina/openapi/core/generators/common/OASModifier.java
index 4ed7aefe2..21ac0b434 100644
--- a/openapi-core/src/main/java/io/ballerina/openapi/core/generators/common/OASModifier.java
+++ b/openapi-core/src/main/java/io/ballerina/openapi/core/generators/common/OASModifier.java
@@ -119,79 +119,98 @@ public OpenAPI modifyWithBallerinaConventions(OpenAPI openapi, Map<String, Strin
         return openapi;
     }
 
-    private static void modifyOASWithObjectPropertyNameInlineSchema(OpenAPI openAPI) {
-        for (Map.Entry<String, PathItem> pathEntry : openAPI.getPaths().entrySet()) {
-            PathItem pathItem = pathEntry.getValue();
-            for (PathItem.HttpMethod method : pathItem.readOperationsMap().keySet()) {
-                Operation operation = pathItem.readOperationsMap().get(method);
-                if (operation.getRequestBody() != null) {
-                    checkForInlineObjects(operation.getRequestBody().getContent());
-                }
-                if (operation.getParameters() != null) {
-                    for (Parameter parameter : operation.getParameters()) {
-                        if (parameter.getSchema() != null) {
-                            updateInlineSchema(parameter.getSchema());
-                        }
-                    }
-                }
-                for (Map.Entry<String, ApiResponse> responseEntry : operation.getResponses().entrySet()) {
-                    ApiResponse response = responseEntry.getValue();
-                    if (response.getContent() != null) {
-                        checkForInlineObjects(response.getContent());
-                    }
-                }
+    private void modifyOASWithObjectPropertyNameInlineSchema(OpenAPI openAPI) {
+        openAPI.getPaths().forEach((path, pathItem) -> processPathItem(pathItem));
+    }
+
+    private void processPathItem(PathItem pathItem) {
+        pathItem.readOperationsMap().forEach((method, operation) -> processOperationWithInlineSchema(operation));
+    }
+
+    private void processOperationWithInlineSchema(Operation operation) {
+        if (operation.getRequestBody() != null) {
+            updateInlineObjectInContent(operation.getRequestBody().getContent());
+        }
+        processParametersWithInlineSchema(operation.getParameters());
+        processResponsesWithInlineSchema(operation.getResponses());
+    }
+
+    private void processParametersWithInlineSchema(List<Parameter> parameters) {
+        if (parameters == null) return;
+
+        for (Parameter parameter : parameters) {
+            if (parameter.getSchema() != null) {
+                updateInlineSchema(parameter.getSchema());
             }
         }
     }
 
+    private void processResponsesWithInlineSchema(Map<String, ApiResponse> responses) {
+        responses.forEach((status, response) -> {
+            if (response.getContent() != null) {
+                updateInlineObjectInContent(response.getContent());
+            }
+        });
+    }
 
-    private static void updateInlineSchema(Schema schema) {
+    private void updateInlineSchema(Schema<?> schema) {
         if (schema == null) {
             return;
         }
-        if (schema instanceof ComposedSchema) {
-            ComposedSchema composedSchema = (ComposedSchema) schema;
-            if (composedSchema.getAllOf() != null) {
-                for (Schema subSchema : composedSchema.getAllOf()) {
-                    updateInlineSchema(subSchema);
-                }
-            }
-            if (composedSchema.getAnyOf() != null) {
-                for (Schema subSchema : composedSchema.getAnyOf()) {
-                    updateInlineSchema(subSchema);
-                }
-            }
-            if (composedSchema.getOneOf() != null) {
-                for (Schema subSchema : composedSchema.getOneOf()) {
-                    updateInlineSchema(subSchema);
-                }
-            }
-        } else if (schema.getType() == null && schema.get$ref() == null && schema.getProperties() != null) {
+        handleInlineSchemaInComposedSchema(schema);
+        handleInlineObjectSchema(schema);
+        handleInlineSchemaItems(schema);
+        handleInlineSchemaInAdditionalProperties(schema);
+        handleInlineSchemaProperties(schema);
+    }
+
+    private void handleInlineSchemaInComposedSchema(Schema<?> schema) {
+        if (schema instanceof ComposedSchema composedSchema) {
+            processSubSchemas(composedSchema.getAllOf());
+            processSubSchemas(composedSchema.getAnyOf());
+            processSubSchemas(composedSchema.getOneOf());
+        }
+    }
+
+    private void processSubSchemas(List<Schema> subSchemas) {
+        if (subSchemas != null) {
+            subSchemas.forEach(this::updateInlineSchema);
+        }
+    }
+
+    private static void handleInlineObjectSchema(Schema<?> schema) {
+        if (isInlineObjectSchema(schema)) {
             Map<String, Schema> properties = schema.getProperties();
-            if (Objects.nonNull(properties) && !properties.isEmpty()) {
-                schema.setProperties(getPropertiesWithBallerinaNameExtension(properties));
-            }
-        } else if (schema instanceof ObjectSchema objectSchema && schema.getProperties() != null) {
-            Map<String, Schema> properties = objectSchema.getProperties();
-            if (Objects.nonNull(properties) && !properties.isEmpty()) {
+            if (properties != null && !properties.isEmpty()) {
                 schema.setProperties(getPropertiesWithBallerinaNameExtension(properties));
             }
         }
+    }
+
+    private static boolean isInlineObjectSchema(Schema<?> schema) {
+        return schema instanceof ObjectSchema ||
+                (schema.getType() == null && schema.get$ref() == null && schema.getProperties() != null);
+    }
+
+    private void handleInlineSchemaItems(Schema<?> schema) {
         if (schema.getItems() != null) {
             updateInlineSchema(schema.getItems());
         }
+    }
+
+    private void handleInlineSchemaInAdditionalProperties(Schema<?> schema) {
         if (schema.getAdditionalProperties() instanceof Schema) {
             updateInlineSchema((Schema) schema.getAdditionalProperties());
         }
+    }
+
+    private void handleInlineSchemaProperties(Schema<?> schema) {
         if (schema.getProperties() != null) {
-            Map<String, Schema> properties = schema.getProperties();
-            for (Map.Entry<String, Schema> property : properties.entrySet()) {
-                updateInlineSchema(property.getValue());
-            }
+            schema.getProperties().values().forEach(this::updateInlineSchema);
         }
     }
 
-    private static void checkForInlineObjects(Map<String, io.swagger.v3.oas.models.media.MediaType> content) {
+    private void updateInlineObjectInContent(Map<String, io.swagger.v3.oas.models.media.MediaType> content) {
         for (Map.Entry<String, io.swagger.v3.oas.models.media.MediaType> entry : content.entrySet()) {
             Schema schema = entry.getValue().getSchema();
             updateInlineSchema(schema);