Fix nested augmentations

[yangtools.git] / data / yang-data-api / src / main / java / org / opendaylight / yangtools / yang / data / api / schema / stream / YangInstanceIdentifierWriter.java

diff --git a/data/yang-data-api/src/main/java/org/opendaylight/yangtools/yang/data/api/schema/stream/YangInstanceIdentifierWriter.java b/data/yang-data-api/src/main/java/org/opendaylight/yangtools/yang/data/api/schema/stream/YangInstanceIdentifierWriter.java
index 03588ede79e7a0a06a483c8652f27055971b8c47..40fad15103bec7beea5098dfe24f80eff356f735 100644 (file)
--- a/data/yang-data-api/src/main/java/org/opendaylight/yangtools/yang/data/api/schema/stream/YangInstanceIdentifierWriter.java
+++ b/data/yang-data-api/src/main/java/org/opendaylight/yangtools/yang/data/api/schema/stream/YangInstanceIdentifierWriter.java
@@ -37,6 +37,7 @@ import org.opendaylight.yangtools.yang.model.api.DataSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.LeafListSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.LeafSchemaNode;
 import org.opendaylight.yangtools.yang.model.api.ListSchemaNode;
+import org.opendaylight.yangtools.yang.model.util.EffectiveAugmentationSchema;
 
 /**
  * Utility for emitting a {@link YangInstanceIdentifier} into a {@link NormalizedNodeStreamWriter} as a set of
@@ -97,14 +98,21 @@ public final class YangInstanceIdentifierWriter implements AutoCloseable {
             final var arg = it.next();
             if (arg instanceof AugmentationIdentifier) {
                 if (!(parent instanceof AugmentationTarget)) {
-                    throw new IOException(parent + " does not support augmentations, cannot resolve" + arg);
+                    throw new IOException(parent + " does not support augmentations, cannot resolve " + arg);
                 }
                 if (reuse) {
                     throw new IOException(parent + " is expecting a nested item, cannot resolve " + arg);
                 }
 
                 final var augId = (AugmentationIdentifier) arg;
-                parent = enterAugmentation((AugmentationTarget) parent, augId);
+                if (parent instanceof DataNodeContainer) {
+                    parent = new EffectiveAugmentationSchema(enterAugmentation((AugmentationTarget) parent, augId),
+                        (DataNodeContainer) parent);
+                } else if (parent instanceof ChoiceSchemaNode) {
+                    throw new IOException(parent + " should not use addressing through " + arg);
+                } else {
+                    throw new IOException("Unhandled parent " + parent + " while resolving " + arg);
+                }
                 writer.startAugmentationNode(augId);
             } else if (arg instanceof NodeWithValue) {
                 if (!(parent instanceof LeafListSchemaNode)) {