Bug 7945: Fix schema validation for augmentation of case

[yangtools.git] / yang / yang-data-impl / src / main / java / org / opendaylight / yangtools / yang / data / impl / schema / tree / CaseEnforcer.java

diff --git a/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/schema/tree/CaseEnforcer.java b/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/schema/tree/CaseEnforcer.java
index 09d061197dc27e5c9ffca9afb963a95a340c8685..548bfb5f9d4fa3b5701fd4edcb2f074e3b57ce50 100644 (file)
--- a/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/schema/tree/CaseEnforcer.java
+++ b/yang/yang-data-impl/src/main/java/org/opendaylight/yangtools/yang/data/impl/schema/tree/CaseEnforcer.java
@@ -10,40 +10,59 @@ package org.opendaylight.yangtools.yang.data.impl.schema.tree;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMap.Builder;
+import com.google.common.collect.Sets;
+
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+
 import org.opendaylight.yangtools.concepts.Immutable;
+import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.AugmentationIdentifier;
 import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.NodeIdentifier;
+import org.opendaylight.yangtools.yang.data.api.YangInstanceIdentifier.PathArgument;
 import org.opendaylight.yangtools.yang.data.api.schema.NormalizedNode;
 import org.opendaylight.yangtools.yang.data.api.schema.tree.DataTreeConfiguration;
 import org.opendaylight.yangtools.yang.data.api.schema.tree.TreeType;
+import org.opendaylight.yangtools.yang.data.impl.schema.SchemaUtils;
+import org.opendaylight.yangtools.yang.model.api.AugmentationSchema;
 import org.opendaylight.yangtools.yang.model.api.ChoiceCaseNode;
 import org.opendaylight.yangtools.yang.model.api.DataSchemaNode;
 
 final class CaseEnforcer implements Immutable {
     private final Map<NodeIdentifier, DataSchemaNode> children;
+    private final Map<AugmentationIdentifier, AugmentationSchema> augmentations;
     private final MandatoryLeafEnforcer enforcer;
 
-    private CaseEnforcer(final Map<NodeIdentifier, DataSchemaNode> children, final MandatoryLeafEnforcer enforcer) {
+    private CaseEnforcer(final Map<NodeIdentifier, DataSchemaNode> children,
+                         final Map<AugmentationIdentifier, AugmentationSchema> augmentations,
+                         final MandatoryLeafEnforcer enforcer) {
         this.children = Preconditions.checkNotNull(children);
+        this.augmentations = Preconditions.checkNotNull(augmentations);
         this.enforcer = Preconditions.checkNotNull(enforcer);
     }
 
     static CaseEnforcer forTree(final ChoiceCaseNode schema, final DataTreeConfiguration treeConfig) {
         final TreeType type = treeConfig.getTreeType();
-        final Builder<NodeIdentifier, DataSchemaNode> builder = ImmutableMap.builder();
+        final Builder<NodeIdentifier, DataSchemaNode> childrenBuilder = ImmutableMap.builder();
+        final Builder<AugmentationIdentifier, AugmentationSchema> augmentationsBuilder = ImmutableMap.builder();
         if (SchemaAwareApplyOperation.belongsToTree(type, schema)) {
             for (final DataSchemaNode child : schema.getChildNodes()) {
                 if (SchemaAwareApplyOperation.belongsToTree(type, child)) {
-                    builder.put(NodeIdentifier.create(child.getQName()), child);
+                    childrenBuilder.put(NodeIdentifier.create(child.getQName()), child);
+                }
+            }
+            for (final AugmentationSchema augment : schema.getAvailableAugmentations()) {
+                if (augment.getChildNodes().stream()
+                        .anyMatch(child -> SchemaAwareApplyOperation.belongsToTree(type, child))) {
+                    augmentationsBuilder.put(SchemaUtils.getNodeIdentifierForAugmentation(augment), augment);
                 }
             }
         }
 
-        final Map<NodeIdentifier, DataSchemaNode> children = builder.build();
-        return children.isEmpty() ? null : new CaseEnforcer(children, MandatoryLeafEnforcer.forContainer(schema,
-                treeConfig));
+        final Map<NodeIdentifier, DataSchemaNode> children = childrenBuilder.build();
+        final Map<AugmentationIdentifier, AugmentationSchema> augmentations = augmentationsBuilder.build();
+        return children.isEmpty() ? null
+                : new CaseEnforcer(children, augmentations, MandatoryLeafEnforcer.forContainer(schema, treeConfig));
     }
 
     Set<Entry<NodeIdentifier, DataSchemaNode>> getChildEntries() {
@@ -54,6 +73,18 @@ final class CaseEnforcer implements Immutable {
         return children.keySet();
     }
 
+    Set<Entry<AugmentationIdentifier, AugmentationSchema>> getAugmentationEntries() {
+        return augmentations.entrySet();
+    }
+
+    Set<AugmentationIdentifier> getAugmentationIdentifiers() {
+        return augmentations.keySet();
+    }
+
+    Set<PathArgument> getAllChildIdentifiers() {
+        return Sets.union(children.keySet(), augmentations.keySet());
+    }
+
     void enforceOnTreeNode(final NormalizedNode<?, ?> normalizedNode) {
         enforcer.enforceOnData(normalizedNode);
     }