- final DataSchemaNode resultFromRecursion = findSchemaNodeOrParentChoiceByName(caseNode, elementName);
- if (resultFromRecursion != null) {
- // this returns top choice node in which child element is found
- return choiceNode;
+ final Deque<Object> resultFromRecursion = findPathToSchemaNodeByName(caseNode, elementName);
+ if (!resultFromRecursion.isEmpty()) {
+ resultFromRecursion.push(choiceNode);
+ if (choiceNode.isAugmenting()) {
+ final AugmentationSchema augment = findCorrespondingAugment(schemaNode, choiceNode);
+ if (augment != null) {
+ resultFromRecursion.push(augment);
+ }
+ }
+ return resultFromRecursion;
+ }
+ }
+ }
+ return result;
+ }
+
+ private static AugmentationSchema findCorrespondingAugment(final DataSchemaNode parent, final DataSchemaNode child) {
+ if (parent instanceof AugmentationTarget && !(parent instanceof ChoiceSchemaNode)) {
+ for (final AugmentationSchema augmentation : ((AugmentationTarget) parent).getAvailableAugmentations()) {
+ final DataSchemaNode childInAugmentation = augmentation.getDataChildByName(child.getQName());
+ if (childInAugmentation != null) {
+ return augmentation;