Merge "Remove unnecessary casts"

[yangtools.git] / yang / yang-data-impl / src / main / java / org / opendaylight / yangtools / yang / data / impl / codec / SchemaTracker.java

diff --git a/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java b/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java
index 79df0ea2006e4528adc52403a66644748ecbaeb7..b4c5a2bd9fd69a071889dc53813b99944b8a9157 100644 (file)
--- a/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java
+++ b/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/codec/SchemaTracker.java
@@ -147,13 +147,7 @@ public final class SchemaTracker {
                 schema = ((NotificationDefinition) parent);
             }
         } else if(parent instanceof ChoiceNode) {
-            for(final ChoiceCaseNode caze : ((ChoiceNode) parent).getCases()) {
-                final DataSchemaNode potential = caze.getDataChildByName(qname);
-                if(potential != null) {
-                    schema = potential;
-                    break;
-                }
-            }
+            schema = findChildInCases((ChoiceNode) parent, qname);
         } else {
             throw new IllegalStateException("Unsupported schema type "+ parent.getClass() +" on stack.");
         }
@@ -161,6 +155,30 @@ public final class SchemaTracker {
         return schema;
     }
 
+    private SchemaNode findChildInCases(final ChoiceNode parent, final QName qname) {
+        DataSchemaNode schema = null;
+        for(final ChoiceCaseNode caze : parent.getCases()) {
+            final DataSchemaNode potential = caze.getDataChildByName(qname);
+            if(potential != null) {
+                schema = potential;
+                break;
+            }
+        }
+        return schema;
+    }
+
+    private SchemaNode findCaseByChild(final ChoiceNode parent, final QName qname) {
+        DataSchemaNode schema = null;
+        for(final ChoiceCaseNode caze : parent.getCases()) {
+            final DataSchemaNode potential = caze.getDataChildByName(qname);
+            if(potential != null) {
+                schema = caze;
+                break;
+            }
+        }
+        return schema;
+    }
+
     public void startList(final PathArgument name) {
         final SchemaNode schema = getSchema(name);
         Preconditions.checkArgument(schema instanceof ListSchemaNode, "Node %s is not a list", schema.getPath());
@@ -218,9 +236,13 @@ public final class SchemaTracker {
 
     public AugmentationSchema startAugmentationNode(final AugmentationIdentifier identifier) {
         LOG.debug("Enter augmentation {}", identifier);
-        final Object parent = getParent();
+        Object parent = getParent();
 
         Preconditions.checkArgument(parent instanceof AugmentationTarget, "Augmentation not allowed under %s", parent);
+        if(parent instanceof ChoiceNode) {
+            final QName name = Iterables.get(identifier.getPossibleChildNames(), 0);
+            parent = findCaseByChild((ChoiceNode) parent, name);
+        }
         Preconditions.checkArgument(parent instanceof DataNodeContainer, "Augmentation allowed only in DataNodeContainer",parent);
         final AugmentationSchema schema = SchemaUtils.findSchemaForAugment((AugmentationTarget) parent, identifier.getPossibleChildNames());
         final HashSet<DataSchemaNode> realChildSchemas = new HashSet<>();