Collect all the updates for subtree writers

So far, when a subtree writer was registered on a list node
and ModificationDiff detected 2 or more updated list items for that writer,
FlatWriterRegistry just picked the first item in list, processed that one
and ignored the rest.

Change-Id: If66db1eaad5a3b5c35e5586f46fd83a0698e1728
Signed-off-by: Maros Marsalek <maros.mars@gmail.com>

3 files changed, 84 insertions, 7 deletions

diff --git a/infra/translate-impl/src/main/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistry.java b/infra/translate-impl/src/main/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistry.java
index 35f1c90a1..67a5da32a 100644
--- a/infra/translate-impl/src/main/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistry.java
+++ b/infra/translate-impl/src/main/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistry.java
@@ -18,6 +18,7 @@ package io.fd.honeycomb.translate.impl.write.registry;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
+import static java.util.stream.Collectors.toMap;
 
 import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableMap;
@@ -33,11 +34,11 @@ import io.fd.honeycomb.translate.write.registry.UpdateFailedException;
 import io.fd.honeycomb.translate.write.registry.WriterRegistry;
 import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -160,25 +161,35 @@ final class FlatWriterRegistry implements WriterRegistry {
                 .orElse(null);
     }
 
-    private Collection<DataObjectUpdate> getParentDataObjectUpdate(final WriteContext ctx,
+    static Collection<DataObjectUpdate> getParentDataObjectUpdate(final WriteContext ctx,
                                                                    final Multimap<InstanceIdentifier<?>, ? extends DataObjectUpdate> updates,
                                                                    final Writer<?> writer) {
         // Now read data for subtree reader root, but first keyed ID is needed and that ID can be cut from updates
-        InstanceIdentifier<?> firstAffectedChildId = ((SubtreeWriter<?>) writer).getHandledChildTypes().stream()
+        return ((SubtreeWriter<?>) writer).getHandledChildTypes().stream()
                 .filter(updates::containsKey)
                 .map(unkeyedId -> updates.get(unkeyedId))
                 .flatMap(doUpdates -> doUpdates.stream())
                 .map(DataObjectUpdate::getId)
-                .findFirst()
-                .get();
+                .map(id -> getSingleParentDataObjectUpdate(ctx, (Multimap<InstanceIdentifier<?>, DataObjectUpdate>) updates, writer, id))
+                // Reduce the list of updates by putting them to a map. If a subtree writer for container gets 2 children updated, it will
+                // get only a single update, however if its a registered on a listand 2 different list items get their children updated,
+                // both updates should be preserved.
+                // Essentially, only group child updates in case the ID from root to writer is identical
+                .collect(toMap(update -> RWUtils.cutId(update.getId(), writer.getManagedDataObjectType()), Function.identity(), (u1, u2) -> u1))
+                .values();
+    }
 
+    private static DataObjectUpdate getSingleParentDataObjectUpdate(WriteContext ctx, Multimap<InstanceIdentifier<?>, DataObjectUpdate> updates, Writer<?> writer, InstanceIdentifier<?> firstAffectedChildId) {
         final InstanceIdentifier<?> parentKeyedId =
                 RWUtils.cutId(firstAffectedChildId, writer.getManagedDataObjectType());
 
         final Optional<? extends DataObject> parentBefore = ctx.readBefore(parentKeyedId);
         final Optional<? extends DataObject> parentAfter = ctx.readAfter(parentKeyedId);
-        return Collections.singleton(
-                DataObjectUpdate.create(parentKeyedId, parentBefore.orNull(), parentAfter.orNull()));
+
+        // Put the parent update data into updates map so that revert can also access the state
+        DataObjectUpdate parentUpdate = DataObjectUpdate.create(parentKeyedId, parentBefore.orNull(), parentAfter.orNull());
+        updates.put(RWUtils.makeIidWildcarded(parentKeyedId), parentUpdate);
+        return parentUpdate;
     }
 
     private void bulkUpdate(
diff --git a/infra/translate-impl/src/test/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistryTest.java b/infra/translate-impl/src/test/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistryTest.java
index bead06ecf..7f4b93e01 100644
--- a/infra/translate-impl/src/test/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistryTest.java
+++ b/infra/translate-impl/src/test/java/io/fd/honeycomb/translate/impl/write/registry/FlatWriterRegistryTest.java
@@ -31,10 +31,13 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.verifyZeroInteractions;
 import static org.mockito.Mockito.when;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableMultimap;
+import com.google.common.collect.Lists;
 import com.google.common.collect.Multimap;
+import com.google.common.collect.Sets;
 import io.fd.honeycomb.translate.util.DataObjects;
 import io.fd.honeycomb.translate.util.DataObjects.DataObject1;
 import io.fd.honeycomb.translate.util.DataObjects.DataObject2;
@@ -44,6 +47,8 @@ import io.fd.honeycomb.translate.write.WriteFailedException;
 import io.fd.honeycomb.translate.write.Writer;
 import io.fd.honeycomb.translate.write.registry.UpdateFailedException;
 import io.fd.honeycomb.translate.write.registry.WriterRegistry;
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -55,6 +60,7 @@ import org.mockito.MockitoAnnotations;
 import org.mockito.stubbing.Answer;
 import org.opendaylight.yangtools.yang.binding.DataObject;
 import org.opendaylight.yangtools.yang.binding.InstanceIdentifier;
+import org.opendaylight.yangtools.yang.binding.KeyedInstanceIdentifier;
 
 public class FlatWriterRegistryTest {
 
@@ -92,6 +98,58 @@ public class FlatWriterRegistryTest {
     }
 
     @Test
+    public void testSubtreeWriterUpdateAggregation() throws Exception {
+        Multimap<InstanceIdentifier<?>, DataObjectUpdate> updates = HashMultimap.create();
+
+        when(ctx.readAfter(DataObject1.IID)).thenReturn(Optional.of(mock(DataObject1.class)));
+        when(ctx.readBefore(DataObject1.IID)).thenReturn(Optional.of(mock(DataObject1.class)));
+
+        Writer<?> writer = SubtreeWriter.createForWriter(Collections.singleton(DataObjects.DataObject1ChildK.IID), writer1);
+
+        InstanceIdentifier<DataObjects.DataObject1ChildK> update1Id = DataObject1.IID.child(DataObjects.DataObject1ChildK.class, new DataObjects.DataObject1ChildKey());
+        InstanceIdentifier<DataObjects.DataObject1ChildK> update2Id = DataObject1.IID.child(DataObjects.DataObject1ChildK.class, new DataObjects.DataObject1ChildKey());
+        updates.putAll(DataObjects.DataObject1ChildK.IID,
+                Lists.newArrayList(
+                        DataObjectUpdate.create(update1Id, mock(DataObjects.DataObject1ChildK.class), mock(DataObjects.DataObject1ChildK.class)),
+                        DataObjectUpdate.create(update2Id, mock(DataObjects.DataObject1ChildK.class), mock(DataObjects.DataObject1ChildK.class))));
+
+        Collection<DataObjectUpdate> parentDataObjectUpdate = FlatWriterRegistry.getParentDataObjectUpdate(ctx, updates, writer);
+        // Just a single update, since there are 2 child updates for a container, they get reduced
+        assertEquals(1, parentDataObjectUpdate.size());
+    }
+
+    @Test
+    public void testSubtreeWriterUpdateAggregationForList() throws Exception {
+        Multimap<InstanceIdentifier<?>, DataObjectUpdate> updates = HashMultimap.create();
+
+        KeyedInstanceIdentifier<DataObjects.DataObject1ChildK, DataObjects.DataObject1ChildKey> parentKeyedId1 =
+                DataObject1.IID.child(DataObjects.DataObject1ChildK.class, new DataObjects.DataObject1ChildKey());
+        KeyedInstanceIdentifier<DataObjects.DataObject1ChildK, DataObjects.DataObject1ChildKey> parentKeyedId2 =
+                DataObject1.IID.child(DataObjects.DataObject1ChildK.class, new DataObjects.DataObject1ChildKey());
+
+        when(ctx.readBefore(parentKeyedId1)).thenReturn(Optional.of(mock(DataObjects.DataObject1ChildK.class)));
+        when(ctx.readAfter(parentKeyedId1)).thenReturn(Optional.of(mock(DataObjects.DataObject1ChildK.class)));
+        when(ctx.readBefore(parentKeyedId2)).thenReturn(Optional.of(mock(DataObjects.DataObject1ChildK.class)));
+        when(ctx.readAfter(parentKeyedId2)).thenReturn(Optional.of(mock(DataObjects.DataObject1ChildK.class)));
+
+        Writer<?> writer = SubtreeWriter.createForWriter(Sets.newHashSet(
+                InstanceIdentifier.create(DataObjects.DataObject1ChildK.class).child(DataObjects.DataObject1ChildK.DataObject1ChildKNested.class),
+                InstanceIdentifier.create(DataObjects.DataObject1ChildK.class).child(DataObjects.DataObject1ChildK.DataObject1ChildKNested2.class)),
+                writer4);
+
+        InstanceIdentifier<DataObjects.DataObject1ChildK.DataObject1ChildKNested> updateList1Id = parentKeyedId1.child(DataObjects.DataObject1ChildK.DataObject1ChildKNested.class);
+        InstanceIdentifier<DataObjects.DataObject1ChildK.DataObject1ChildKNested> updateList2Id = parentKeyedId2.child(DataObjects.DataObject1ChildK.DataObject1ChildKNested.class);
+        updates.putAll(DataObjects.DataObject1ChildK.DataObject1ChildKNested.IID,
+                Lists.newArrayList(
+                        DataObjectUpdate.create(updateList1Id, mock(DataObjects.DataObject1ChildK.DataObject1ChildKNested.class), mock(DataObjects.DataObject1ChildK.DataObject1ChildKNested.class)),
+                        DataObjectUpdate.create(updateList2Id, mock(DataObjects.DataObject1ChildK.DataObject1ChildKNested.class), mock(DataObjects.DataObject1ChildK.DataObject1ChildKNested.class))));
+
+        Collection<DataObjectUpdate> parentDataObjectUpdate = FlatWriterRegistry.getParentDataObjectUpdate(ctx, updates, writer);
+        // 2 updates for 2 different list items
+        assertEquals(2, parentDataObjectUpdate.size());
+    }
+
+    @Test
     public void testMultipleUpdatesForSingleWriter() throws Exception {
         final FlatWriterRegistry flatWriterRegistry =
                 new FlatWriterRegistry(ImmutableMap.of(DataObject1.IID, writer1, DataObject2.IID, writer2));
diff --git a/infra/translate-utils/src/test/java/io/fd/honeycomb/translate/util/DataObjects.java b/infra/translate-utils/src/test/java/io/fd/honeycomb/translate/util/DataObjects.java
index 77db65c5d..f3da1c767 100644
--- a/infra/translate-utils/src/test/java/io/fd/honeycomb/translate/util/DataObjects.java
+++ b/infra/translate-utils/src/test/java/io/fd/honeycomb/translate/util/DataObjects.java
@@ -69,6 +69,14 @@ public class DataObjects {
                 RWUtils.makeIidWildcarded(InstanceIdentifier.create(DataObject1.class).child(DataObject1ChildK.class));
         InstanceIdentifier<DataObject1ChildK> INTERNALLY_KEYED_IID = InstanceIdentifier.create(DataObject1.class)
                 .child(DataObject1ChildK.class, new DataObject1ChildKey());
+
+        public interface DataObject1ChildKNested extends DataObject, ChildOf<DataObject1ChildK> {
+            InstanceIdentifier<DataObject1ChildK.DataObject1ChildKNested> IID = DataObject1ChildK.IID.child(DataObject1ChildK.DataObject1ChildKNested.class);
+        }
+
+        public interface DataObject1ChildKNested2 extends DataObject, ChildOf<DataObject1ChildK> {
+            InstanceIdentifier<DataObject1ChildK.DataObject1ChildKNested2> IID = DataObject1ChildK.IID.child(DataObject1ChildK.DataObject1ChildKNested2.class);
+        }
     }
 
     public static class DataObject1ChildKey implements Identifier<DataObject1ChildK> {