Fix ref handling in Discriminator mappings

Author: andurairaj

modules/swagger-parser-v3/src/main/java/io/swagger/v3/parser/processors/ExternalRefProcessor.java

@@ -4,10 +4,15 @@
 import java.net.URI;
 import java.nio.file.Paths;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Optional;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import io.swagger.v3.oas.models.Components;
 import io.swagger.v3.oas.models.OpenAPI;
@@ -246,9 +251,31 @@ private void processSchema(Schema property, String file) {
             }
             if (property instanceof ComposedSchema) {
                 ComposedSchema composed = (ComposedSchema) property;
+                final Map<String, String> refMap = Optional.ofNullable(composed.getDiscriminator())
+                                .map(Discriminator::getMapping).orElse(Collections.emptyMap()).entrySet()
+                                .stream().collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
+                Map<String, Schema> refCache = (!refMap.isEmpty() &&
+                        (composed.getAnyOf() != null || composed.getOneOf() != null)) ? Stream.of(
+                                composed.getAnyOf(), composed.getOneOf()
+                        )
+                            .filter(Objects::nonNull)
+                            .filter(l -> !l.isEmpty())
+                            .flatMap(Collection::stream)
+                            .filter(s -> s.get$ref() != null)
+                            .collect(Collectors.toMap(Schema::get$ref, Function.identity())) : Collections.emptyMap();
+
                 processProperties(composed.getAllOf(), file);
                 processProperties(composed.getAnyOf(), file);
                 processProperties(composed.getOneOf(), file);
+
+                if (!refMap.isEmpty() && !refCache.isEmpty()) {
+                    refCache.entrySet()
+                            .stream().filter(e -> !e.getKey().equals(e.getValue().get$ref()))
+                            .forEach(entry -> {
+                                String newRef = entry.getValue().get$ref();
+                                property.getDiscriminator().getMapping().put(refMap.get(entry.getKey()), newRef);
+                            });
+                }
             }
         }
     }

modules/swagger-parser-v3/src/main/java/io/swagger/v3/parser/processors/PathsProcessor.java

@@ -256,15 +256,33 @@ else if(model instanceof ComposedSchema) {
                 for (Schema innerModel : composedSchema.getAllOf()) {
                     updateRefs(innerModel, pathRef);
                 }
-            }if (composedSchema.getAnyOf() != null) {
-                for(Schema innerModel : composedSchema.getAnyOf()) {
-                    updateRefs(innerModel, pathRef);
-                }
-            }if (composedSchema.getOneOf() != null) {
-                for (Schema innerModel : composedSchema.getOneOf()) {
-                    updateRefs(innerModel, pathRef);
+            }
+            if (composedSchema.getAnyOf() != null || composedSchema.getOneOf() != null) {
+                // Map to cache old - new refs in composed schemas
+                Map<String, String> refMappings = composedSchema.getDiscriminator() != null &&
+                        composedSchema.getDiscriminator().getMapping() != null ? new HashMap<>() : null;
+
+                Stream.of(composedSchema.getAnyOf(), composedSchema.getOneOf())
+                        .filter(Objects::nonNull).filter(l -> !l.isEmpty())
+                        .flatMap(Collection::stream)
+                        .forEach(innerModel -> {
+                            String oldRef = innerModel.get$ref();
+                            updateRefs(innerModel, pathRef);
+                            if(oldRef != null && refMappings != null && !oldRef.equals(innerModel.get$ref())) {
+                                refMappings.put(oldRef, innerModel.get$ref());
+                            }
+                });
+                // Update refs in discriminator mappings
+                if(refMappings != null && !refMappings.isEmpty()) {
+                    Map<String, String> discriminatorMappings = composedSchema.getDiscriminator().getMapping();
+                    for(String key : discriminatorMappings.keySet()) {
+                        if(refMappings.containsKey(discriminatorMappings.get(key))) {
+                            discriminatorMappings.put(key, refMappings.get(discriminatorMappings.get(key)));
+                        }
+                    }
                 }
             }
+
         }
         else if(model instanceof ArraySchema) {
             ArraySchema arraySchema = (ArraySchema) model;

modules/swagger-parser-v3/src/main/java/io/swagger/v3/parser/processors/SchemaProcessor.java

@@ -14,9 +14,12 @@
 import io.swagger.v3.parser.util.OpenAPIDeserializer;
 import io.swagger.v3.parser.util.RefUtils;
 
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Stream;
 
 import static io.swagger.v3.parser.util.RefUtils.computeRefFormat;
 import static io.swagger.v3.parser.util.RefUtils.isAnExternalRefFormat;
@@ -157,33 +160,21 @@ public void processComposedSchema(ComposedSchema composedSchema) {
                     }
                 }
             }
-        }if(composedSchema.getOneOf() != null){
-            final List<Schema> schemas = composedSchema.getOneOf();
-            if (schemas != null) {
-                for (Schema schema : schemas) {
-                    if (schema.get$ref() != null) {
-                        String oldRef = schema.get$ref();
-                        processReferenceSchema(schema);
-                        String newRef = schema.get$ref();
-                        changeDiscriminatorMapping(composedSchema, oldRef, newRef);
-                    } else {
-                        processSchemaType(schema);
-                    }
-                }
-            }
-        }if(composedSchema.getAnyOf() != null){
-            final List<Schema> schemas = composedSchema.getAnyOf();
-            if (schemas != null) {
-                for (Schema schema : schemas) {
-                    if (schema.get$ref() != null) {
-                        processReferenceSchema(schema);
-                    } else {
-                        processSchemaType(schema);
-                    }
-                }
-            }
         }
-
+        if(composedSchema.getOneOf() != null || composedSchema.getAnyOf() != null) {
+            Stream.of(composedSchema.getOneOf(), composedSchema.getAnyOf())
+                    .filter(Objects::nonNull).filter(l -> !l.isEmpty()).flatMap(Collection::stream)
+                    .forEach(schema -> {
+                        if (schema.get$ref() != null) {
+                            String oldRef = schema.get$ref();
+                            processReferenceSchema(schema);
+                            String newRef = schema.get$ref();
+                            changeDiscriminatorMapping(composedSchema, oldRef, newRef);
+                        } else {
+                            processSchemaType(schema);
+                        }
+                    });
+        }
     }
 
     private void changeDiscriminatorMapping(ComposedSchema composedSchema, String oldRef, String newRef) {