Fix support for nested augmentations when augmenting lists

Change-Id: I96e7db8f295c9c3d5b14395c7785574d12d76ea9
Signed-off-by: Marek Gradzki <mgradzki@cisco.com>

3 files changed, 104 insertions, 8 deletions

diff --git a/infra/data-impl/src/main/java/io/fd/honeycomb/data/impl/ModificationDiff.java b/infra/data-impl/src/main/java/io/fd/honeycomb/data/impl/ModificationDiff.java
index 9903d99fc..863f8abb2 100644
--- a/infra/data-impl/src/main/java/io/fd/honeycomb/data/impl/ModificationDiff.java
+++ b/infra/data-impl/src/main/java/io/fd/honeycomb/data/impl/ModificationDiff.java
@@ -299,15 +299,26 @@ final class ModificationDiff {
         // yangtools. The hierarchy does not use e.g. SchemaNode class for all types
         private final Object parentNode;
         private final Object schemaNode;
+        private final boolean updateParentNode;
 
-        Modification(final YangInstanceIdentifier id,
-                     final DataTreeCandidateNode dataCandidate,
-                     final Object parentNode,
-                     final Object schemaNode) {
+        private Modification(final YangInstanceIdentifier id,
+                             final DataTreeCandidateNode dataCandidate,
+                             final Object parentNode,
+                             final Object schemaNode,
+                             final boolean updateParentNode) {
             this.id = id;
             this.dataCandidate = dataCandidate;
             this.parentNode = parentNode;
             this.schemaNode = schemaNode;
+            // controls process of updating parent node while moving down the schema tree:
+            this.updateParentNode = updateParentNode;
+        }
+
+        Modification(final YangInstanceIdentifier id,
+                     final DataTreeCandidateNode dataCandidate,
+                     final Object parentNode,
+                     final Object schemaNode) {
+            this(id, dataCandidate, parentNode, schemaNode, true);
         }
 
         Modification(final YangInstanceIdentifier id,
@@ -397,12 +408,19 @@ final class ModificationDiff {
                 .map(child -> {
                     final YangInstanceIdentifier childId = id.node(child.getIdentifier());
                     final Object schemaChild = schemaChild(schemaNode, child.getIdentifier());
+
                     // An augment cannot change other augment, so we do not update parent node if we are streaming
                     // children of AugmentationSchema (otherwise we would fail to find schema for nested augmentations):
-                    final Object newParent = (schemaNode instanceof AugmentationSchema)
-                        ? parentNode
-                        : schemaNode;
-                    return new Modification(childId, child, newParent, schemaChild);
+                    if (updateParentNode) {
+                        if (schemaNode instanceof AugmentationSchema) {
+                            // child nodes would not have nested augmentations, so we stop moving parentNode:
+                            return new Modification(childId, child, parentNode, schemaChild, false);
+                        } else {
+                            // update parent node:
+                            return new Modification(childId, child, schemaNode, schemaChild, true);
+                        }
+                    }
+                    return new Modification(childId, child, parentNode, schemaChild, updateParentNode);
                 });
         }
 
diff --git a/infra/it/it-test/src/test/java/io/fd/honeycomb/data/impl/NestedAugmentationWriteTest.java b/infra/it/it-test/src/test/java/io/fd/honeycomb/data/impl/NestedAugmentationWriteTest.java
index ef7430dd3..351b88e4a 100644
--- a/infra/it/it-test/src/test/java/io/fd/honeycomb/data/impl/NestedAugmentationWriteTest.java
+++ b/infra/it/it-test/src/test/java/io/fd/honeycomb/data/impl/NestedAugmentationWriteTest.java
@@ -29,6 +29,7 @@ import io.fd.honeycomb.translate.write.WriteContext;
 import io.fd.honeycomb.translate.write.Writer;
 import io.fd.honeycomb.translate.write.registry.WriterRegistry;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import org.junit.Test;
@@ -41,12 +42,19 @@ import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.t
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.FromAugmentAugmentBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.FromAugmentListAugment;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.FromAugmentListAugmentBuilder;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.ListAugment;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.ListAugmentBuilder;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.ListFromAugmentAugment;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.ListFromAugmentAugmentBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.SimpleAugment;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.SimpleAugmentBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.SimpleNestedAugment;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.SimpleNestedAugmentBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.FromAugment;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.FromAugmentBuilder;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.ListFromAugment;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.ListFromAugmentBuilder;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.ListFromAugmentKey;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.from.augment.FromAugment2;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.from.augment.FromAugment2Builder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.hc.aug.test.rev161222.aug.target.from.augment.FromAugmentEntry;
@@ -73,6 +81,12 @@ public class NestedAugmentationWriteTest extends AbstractInfraTest {
         AUG_TARGET_ID.augmentation(FromAugmentAugment.class);
     private static final InstanceIdentifier<FromAugment> FROM_AUGMENT_ID =
         FROM_AUGMENT_AUGMENT_ID.child(FromAugment.class);
+    private static final InstanceIdentifier<ListAugment> LIST_AUGMENT_ID =
+        AUG_TARGET_ID.augmentation(ListAugment.class);
+    private static final InstanceIdentifier<ListFromAugment> LIST_FROM_AUGMENT_ID =
+        LIST_AUGMENT_ID.child(ListFromAugment.class);
+    private static final InstanceIdentifier<ListFromAugmentAugment> LIST_FROM_AUGMENT_AUGMENT_ID =
+        LIST_FROM_AUGMENT_ID.augmentation(ListFromAugmentAugment.class);
     private static final InstanceIdentifier<SimpleAugment> SIMPLE_AUGMENT_ID =
         AUG_TARGET_ID.augmentation(SimpleAugment.class);
     private static final InstanceIdentifier<FromAugment2Augment> FROM_AUGMENT2_AUGMENT_ID =
@@ -91,6 +105,9 @@ public class NestedAugmentationWriteTest extends AbstractInfraTest {
 
     private final Writer<AugTarget> augTargetWriter = mockWriter(AUG_TARGET_ID);
     private final Writer<FromAugment> fromAugmentWriter = mockWriter(FROM_AUGMENT_ID);
+    private final Writer<ListFromAugment> listFromAugmentWriter = mockWriter(LIST_FROM_AUGMENT_ID);
+    private final Writer<ListFromAugmentAugment> listFromAugmentAugmentWriter =
+        mockWriter(LIST_FROM_AUGMENT_AUGMENT_ID);
     private final Writer<FromAugment2> fromAugment2Writer = mockWriter(FROM_AUGMENT2_ID);
     private final Writer<FromAugmentEntry> fromAugmentListWriter = mockWriter(FROM_AUGMENT_ENTRY_ID);
 
@@ -118,6 +135,8 @@ public class NestedAugmentationWriteTest extends AbstractInfraTest {
         writerRegistry = new FlatWriterRegistryBuilder()
             .add(augTargetWriter)
             .add(fromAugmentWriter)
+            .add(listFromAugmentWriter)
+            .add(listFromAugmentAugmentWriter)
             .add(simpleAugmentWriter)
             .add(fromAugment2Writer)
             .add(simpleNestedAugmentWriter)
@@ -226,6 +245,47 @@ public class NestedAugmentationWriteTest extends AbstractInfraTest {
             .update(eq(keyedNestedList2), eq(null), eq(entries.get(1)), any(WriteContext.class));
     }
 
+    @Test
+    public void testListNestedAugmentationWrite() throws Exception {
+        // tests augmenting list that already comes from augment
+        final ModifiableDataTreeDelegator modifiableDataTreeDelegator =
+            new ModifiableDataTreeDelegator(serializer, dataTree, schemaContext, writerRegistry, contextBroker);
+
+        final DataModification dataModification = modifiableDataTreeDelegator.newModification();
+
+        final ListFromAugmentAugment listAugmentation = new ListFromAugmentAugmentBuilder()
+            .setNewLeaf("new-leaf-val").build();
+        final List<ListFromAugment> list = Collections.singletonList(
+            new ListFromAugmentBuilder()
+                .setSomeLeaf("some-leaf-val")
+                .addAugmentation(ListFromAugmentAugment.class, listAugmentation).build());
+        final AugTarget data = new AugTargetBuilder()
+            .setSomeLeaf("aug-target-leaf-val")
+            .addAugmentation(ListAugment.class, new ListAugmentBuilder().setListFromAugment(list).build()).build();
+
+        final Map.Entry<YangInstanceIdentifier, NormalizedNode<?, ?>> normalizedNode =
+            serializer.toNormalizedNode(AUG_TARGET_ID, data);
+        dataModification.write(normalizedNode.getKey(), normalizedNode.getValue());
+        dataModification.commit();
+
+        // verify aug target update:
+        verify(augTargetWriter).update(eq(AUG_TARGET_ID), eq(null), eq(data), any(WriteContext.class));
+
+        // verify list customizer update:
+        final KeyedInstanceIdentifier<ListFromAugment, ListFromAugmentKey> keyedNestedList =
+            LIST_AUGMENT_ID.child(ListFromAugment.class, new ListFromAugmentKey("some-leaf-val"));
+        final ArgumentCaptor<DataObject> doCaptor = ArgumentCaptor.forClass(DataObject.class);
+        verify(listFromAugmentWriter)
+            .update(eq(keyedNestedList), eq(null), doCaptor.capture(), any(WriteContext.class));
+        assertEquals(list.get(0).getSomeLeaf(), ((ListFromAugment) doCaptor.getValue()).getSomeLeaf());
+
+        // verify list augmentation customizer update:
+        verify(listFromAugmentAugmentWriter)
+            .update(eq(keyedNestedList.augmentation(ListFromAugmentAugment.class)), eq(null), doCaptor.capture(),
+                any(WriteContext.class));
+        assertEquals(listAugmentation.getNewLeaf(), ((ListFromAugmentAugment) doCaptor.getValue()).getNewLeaf());
+    }
+
     private AugTarget augTarget(FromAugment fromAugment) {
         return new AugTargetBuilder()
             .setSomeLeaf("aug-target-leaf-val")
diff --git a/infra/it/test-model/src/main/yang/hc-aug-test.yang b/infra/it/test-model/src/main/yang/hc-aug-test.yang
index 8a83ab56f..95179d422 100644
--- a/infra/it/test-model/src/main/yang/hc-aug-test.yang
+++ b/infra/it/test-model/src/main/yang/hc-aug-test.yang
@@ -27,6 +27,17 @@ module hc-aug-test {
   }
 
   augment "/aug-target" {
+    ext:augment-identifier "list-augment";
+    list list-from-augment {
+      key some-leaf;
+
+      leaf some-leaf {
+        type string;
+      }
+    }
+  }
+
+  augment "/aug-target" {
     ext:augment-identifier "simple-augment";
     leaf simple-augment-leaf {
         type string;
@@ -60,4 +71,11 @@ module hc-aug-test {
     }
   }
 
+  augment "/aug-target/list-from-augment" {
+    ext:augment-identifier "list-from-augment-augment";
+    leaf new-leaf {
+        type string;
+    }
+  }
+
 }