Cleanup deprecation warnings in bgp/openconfig-state

Eliminate some leftovers from codegen change.

Change-Id: I232ac497b2007a8a916b00ca6fe623e8d4f6e2fb
Signed-off-by: Robert Varga <robert.varga@pantheon.tech>

diff --git a/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/GlobalUtil.java b/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/GlobalUtil.java
index db125b4676662d85c17b1763a67a5220b2fd5f94..4d68e07c88cf1a00175a22887897d42b25fa9e63 100644 (file)
--- a/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/GlobalUtil.java
+++ b/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/GlobalUtil.java
@@ -15,13 +15,11 @@ import org.opendaylight.protocol.bgp.openconfig.spi.BGPTableTypeRegistryConsumer
 import org.opendaylight.protocol.bgp.rib.spi.state.BGPRibState;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi.safi.list.AfiSafi;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi.safi.list.AfiSafiBuilder;
-import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi.safi.list.afi.safi.State;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi.safi.list.afi.safi.StateBuilder;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.global.base.AfiSafisBuilder;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.top.bgp.Global;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.top.bgp.GlobalBuilder;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.types.rev151009.AfiSafiType;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.GlobalAfiSafiStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.GlobalAfiSafiStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.rib.rev180329.rib.TablesKey;
 
@@ -88,11 +86,14 @@ public final class GlobalUtil {
         if (!optAfiSafi.isPresent()) {
             return null;
         }
-        final State state = new StateBuilder()
-                .addAugmentation(GlobalAfiSafiStateAugmentation.class, new GlobalAfiSafiStateAugmentationBuilder()
+        return new AfiSafiBuilder()
+                .setAfiSafiName(optAfiSafi.get())
+                .setState(new StateBuilder()
+                    .addAugmentation(new GlobalAfiSafiStateAugmentationBuilder()
                         .setTotalPaths(ribState.getPathCount(tablesKey))
                         .setTotalPrefixes(ribState.getPrefixesCount(tablesKey))
-                        .build()).build();
-        return new AfiSafiBuilder().setAfiSafiName(optAfiSafi.get()).setState(state).build();
+                        .build())
+                    .build())
+                .build();
     }
 }

diff --git a/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/NeighborUtil.java b/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/NeighborUtil.java
index 58435d75bbe58b3b0925bb44deed6c352ff2d179..4d33a6e5c3d7950d68033d43ebc7aa06e0b1ed38 100644 (file)
--- a/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/NeighborUtil.java
+++ b/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/NeighborUtil.java
@@ -66,7 +66,6 @@ import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.open
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.BgpNeighborStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiGracefulRestartStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiGracefulRestartStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborErrorHandlingStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborErrorHandlingStateAugmentationBuilder;
@@ -74,9 +73,7 @@ import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.open
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborGracefulRestartStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTimersStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTimersStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTransportStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTransportStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.network.instances.network.instance.protocols.protocol.bgp.neighbors.neighbor.state.MessagesBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.network.instances.network.instance.protocols.protocol.bgp.neighbors.neighbor.state.messages.Received;
@@ -159,10 +156,10 @@ public final class NeighborUtil {
         }
         final StateBuilder builder = new StateBuilder();
         if (sessionState != null) {
-            builder.addAugmentation(NeighborStateAugmentation.class, buildCapabilityState(sessionState));
+            builder.addAugmentation(buildCapabilityState(sessionState));
         }
         if (bgpPeerMessagesState != null) {
-            builder.addAugmentation(BgpNeighborStateAugmentation.class, buildMessageState(bgpPeerMessagesState));
+            builder.addAugmentation(buildMessageState(bgpPeerMessagesState));
         }
         return builder.build();
     }
@@ -180,13 +177,15 @@ public final class NeighborUtil {
         // convert neighbor uptime which is in milliseconds to time-ticks which is
         // hundredth of a second, and handle roll-over scenario
         final long uptimeTicks = neighbor.getUpTime() / 10 % TIMETICK_ROLLOVER_VALUE;
-        final NeighborTimersStateAugmentation timerState = new NeighborTimersStateAugmentationBuilder()
-                .setNegotiatedHoldTime(BigDecimal.valueOf(neighbor.getNegotiatedHoldTime()))
-                .setUptime(new Timeticks(Uint32.valueOf(uptimeTicks))).build();
 
-        return new TimersBuilder().setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp
-                .neighbor.group.timers.StateBuilder()
-                .addAugmentation(NeighborTimersStateAugmentation.class, timerState).build()).build();
+        return new TimersBuilder()
+                .setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group
+                    .timers.StateBuilder()
+                        .addAugmentation(new NeighborTimersStateAugmentationBuilder()
+                            .setNegotiatedHoldTime(BigDecimal.valueOf(neighbor.getNegotiatedHoldTime()))
+                            .setUptime(new Timeticks(Uint32.valueOf(uptimeTicks))).build())
+                        .build())
+                .build();
     }
 
     /**
@@ -199,13 +198,16 @@ public final class NeighborUtil {
         if (neighbor == null) {
             return null;
         }
-        final NeighborTransportStateAugmentation transportState = new NeighborTransportStateAugmentationBuilder()
-                .setLocalPort(neighbor.getLocalPort()).setRemoteAddress(convertIpAddress(neighbor.getRemoteAddress()))
-                .setRemotePort(neighbor.getRemotePort()).build();
 
         return new TransportBuilder().setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009
                 .bgp.neighbor.group.transport.StateBuilder()
-                .addAugmentation(NeighborTransportStateAugmentation.class, transportState).build()).build();
+                    .addAugmentation(new NeighborTransportStateAugmentationBuilder()
+                        .setLocalPort(neighbor.getLocalPort())
+                        .setRemoteAddress(convertIpAddress(neighbor.getRemoteAddress()))
+                        .setRemotePort(neighbor.getRemotePort())
+                        .build())
+                    .build())
+                .build();
     }
 
     /**
@@ -218,10 +220,12 @@ public final class NeighborUtil {
         if (errorHandlingState == null) {
             return null;
         }
-        return new ErrorHandlingBuilder().setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp
-                .rev151009.bgp.neighbor.group.error.handling.StateBuilder()
-                .addAugmentation(NeighborErrorHandlingStateAugmentation.class,
-                        buildErrorHandlingState(errorHandlingState.getErroneousUpdateReceivedCount())).build()).build();
+        return new ErrorHandlingBuilder()
+                .setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group
+                    .error.handling.StateBuilder()
+                        .addAugmentation(buildErrorHandlingState(errorHandlingState.getErroneousUpdateReceivedCount()))
+                        .build())
+                .build();
     }
 
     /**
@@ -240,7 +244,7 @@ public final class NeighborUtil {
 
         return new GracefulRestartBuilder().setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp
                 .rev151009.bgp.graceful.restart.graceful.restart.StateBuilder()
-                .addAugmentation(NeighborGracefulRestartStateAugmentation.class, gracefulRestartState).build()).build();
+                .addAugmentation(gracefulRestartState).build()).build();
     }
 
     /**
@@ -358,8 +362,7 @@ public final class NeighborUtil {
                     .setSent(neighbor.getPrefixesSentCount(tablesKey)).build());
         }
         return new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi
-                .safi.list.afi.safi.StateBuilder().addAugmentation(NeighborAfiSafiStateAugmentation.class,
-                builder.build()).build();
+                .safi.list.afi.safi.StateBuilder().addAugmentation(builder.build()).build();
     }
 
     private static org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi
@@ -375,7 +378,7 @@ public final class NeighborUtil {
         return new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi
                 .safi.list.afi.safi.GracefulRestartBuilder().setState(new org.opendaylight.yang.gen.v1.http.openconfig
                 .net.yang.bgp.multiprotocol.rev151009.bgp.common.afi.safi.list.afi.safi.graceful.restart.StateBuilder()
-                .addAugmentation(NeighborAfiSafiGracefulRestartStateAugmentation.class, builder).build()).build();
+                .addAugmentation(builder).build()).build();
     }
 
     /**

diff --git a/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/PeerGroupUtil.java b/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/PeerGroupUtil.java
index d66a18b79ad907fe61191620d1670cc315d144be..10ecda237c767f2406ad54a0c99b98531b5f068b 100644 (file)
--- a/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/PeerGroupUtil.java
+++ b/bgp/openconfig-state/src/main/java/org/opendaylight/protocol/bgp/state/PeerGroupUtil.java
@@ -7,8 +7,10 @@
  */
 package org.opendaylight.protocol.bgp.state;
 
+import com.google.common.collect.ImmutableMap;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import org.eclipse.jdt.annotation.NonNull;
 import org.eclipse.jdt.annotation.Nullable;
@@ -43,10 +45,11 @@ public final class PeerGroupUtil {
             return null;
         }
 
-        final List<PeerGroup> peerGroupsList = peerGroups.entrySet().stream()
-                .map(entry -> buildPeerGroupState(entry.getKey(), entry.getValue()))
-                .collect(Collectors.toList());
-        return new PeerGroupsBuilder().setPeerGroup(peerGroupsList).build();
+        return new PeerGroupsBuilder()
+                .setPeerGroup(peerGroups.entrySet().stream()
+                    .map(entry -> buildPeerGroupState(entry.getKey(), entry.getValue()))
+                    .collect(ImmutableMap.toImmutableMap(PeerGroup::key, Function.identity())))
+                .build();
     }
 
     /**
@@ -66,6 +69,6 @@ public final class PeerGroupUtil {
         return new PeerGroupBuilder()
                 .setPeerGroupName(groupId)
                 .setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group
-                        .StateBuilder().addAugmentation(PeerGroupStateAugmentation.class, groupState).build()).build();
+                        .StateBuilder().addAugmentation(groupState).build()).build();
     }
 }

diff --git a/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/NeighborUtilTest.java b/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/NeighborUtilTest.java
index 9b48b86c89fae77d108a9d06fdf3b22d66b65986..cbc7094bc0d0a2f9bc714db6ea1289b41dcf296b 100644 (file)
--- a/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/NeighborUtilTest.java
+++ b/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/NeighborUtilTest.java
@@ -38,9 +38,7 @@ import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.n
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.types.rev151009.AfiSafiType;
 import org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.types.rev151009.IPV4UNICAST;
 import org.opendaylight.yang.gen.v1.urn.ietf.params.xml.ns.yang.ietf.yang.types.rev130715.Timeticks;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiGracefulRestartStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiGracefulRestartStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTimersStateAugmentation;
@@ -107,7 +105,7 @@ public class NeighborUtilTest {
                 .setNegotiatedHoldTime(BigDecimal.valueOf(90L)).setUptime(new Timeticks(Uint32.valueOf(500))).build();
         final Timers expectedTimers = new TimersBuilder().setState(
                 new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group.timers
-                .StateBuilder().addAugmentation(NeighborTimersStateAugmentation.class, timerStateAug).build())
+                .StateBuilder().addAugmentation(timerStateAug).build())
                 .build();
         assertEquals(expectedTimers, NeighborUtil.buildTimer(timerState));
     }
@@ -122,8 +120,7 @@ public class NeighborUtilTest {
                 .setNegotiatedHoldTime(BigDecimal.valueOf(90L)).setUptime(new Timeticks(Uint32.valueOf(5))).build();
         final Timers expectedTimers = new TimersBuilder().setState(
                 new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group.timers
-                .StateBuilder().addAugmentation(NeighborTimersStateAugmentation.class, timerStateAug).build())
-                .build();
+                .StateBuilder().addAugmentation(timerStateAug).build()).build();
         assertEquals(expectedTimers, NeighborUtil.buildTimer(timerState));
     }
 
@@ -143,7 +140,7 @@ public class NeighborUtilTest {
                 NeighborUtil.buildAfisSafisState(this.bgpAfiSafiState, this.tableRegistry));
 
         final GracefulRestart graceful = new GracefulRestartBuilder()
-                .setState(new StateBuilder().addAugmentation(NeighborAfiSafiGracefulRestartStateAugmentation.class,
+                .setState(new StateBuilder().addAugmentation(
                         new NeighborAfiSafiGracefulRestartStateAugmentationBuilder()
                                 .setAdvertised(false)
                                 .setReceived(false)
@@ -155,8 +152,7 @@ public class NeighborUtilTest {
         final org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp.common.afi.safi
                 .list.afi.safi.State afiSafiState = new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp
                 .multiprotocol.rev151009.bgp.common.afi.safi.list.afi.safi.StateBuilder()
-                .addAugmentation(NeighborAfiSafiStateAugmentation.class, new NeighborAfiSafiStateAugmentationBuilder()
-                        .setActive(false).build()).build();
+                .addAugmentation(new NeighborAfiSafiStateAugmentationBuilder().setActive(false).build()).build();
 
         this.afiSafi = Optional.of(IPV4UNICAST.class);
         final AfiSafi expected = new AfiSafiBuilder().setAfiSafiName(this.afiSafi.get())

diff --git a/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/StateProviderImplTest.java b/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/StateProviderImplTest.java
index dcfd59d66dab1bd931f63b0d18127c9eba16302a..6420661e53757c5375b85bb074963db44337b3d3 100644 (file)
--- a/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/StateProviderImplTest.java
+++ b/bgp/openconfig-state/src/test/java/org/opendaylight/protocol/bgp/state/StateProviderImplTest.java
@@ -33,6 +33,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ScheduledExecutorService;
@@ -116,23 +117,15 @@ import org.opendaylight.yang.gen.v1.urn.ietf.params.xml.ns.yang.ietf.inet.types.
 import org.opendaylight.yang.gen.v1.urn.ietf.params.xml.ns.yang.ietf.yang.types.rev130715.Timeticks;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.BgpNeighborStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.BgpNeighborStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.GlobalAfiSafiStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.GlobalAfiSafiStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiGracefulRestartStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiGracefulRestartStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborAfiSafiStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborErrorHandlingStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborErrorHandlingStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborGracefulRestartStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborGracefulRestartStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborStateAugmentation;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTimersStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTimersStateAugmentationBuilder;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTransportStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NeighborTransportStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.NetworkInstanceProtocol;
-import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.PeerGroupStateAugmentation;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.PeerGroupStateAugmentationBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.network.instances.network.instance.protocols.protocol.bgp.neighbors.neighbor.state.MessagesBuilder;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.bgp.openconfig.extensions.rev180329.network.instances.network.instance.protocols.protocol.bgp.neighbors.neighbor.state.messages.ReceivedBuilder;
@@ -494,59 +487,62 @@ public class StateProviderImplTest extends AbstractDataBrokerTest {
                     .build());
         final AfiSafi afiSafi = new AfiSafiBuilder()
                 .setAfiSafiName(IPV4UNICAST.class)
-                .setGracefulRestart(new GracefulRestartBuilder().setState(new StateBuilder().setEnabled(false)
-                        .addAugmentation(NeighborAfiSafiGracefulRestartStateAugmentation.class,
-                                new NeighborAfiSafiGracefulRestartStateAugmentationBuilder()
-                                        .setAdvertised(true)
-                                        .setReceived(true)
-                                        .setLlStaleTimer(Uint32.valueOf(60))
-                                        .setLlAdvertised(true)
-                                        .setLlReceived(true)
-                                        .build())
-                        .build()).build())
+                .setGracefulRestart(new GracefulRestartBuilder()
+                    .setState(new StateBuilder().setEnabled(false)
+                        .addAugmentation(new NeighborAfiSafiGracefulRestartStateAugmentationBuilder()
+                            .setAdvertised(true)
+                            .setReceived(true)
+                            .setLlStaleTimer(Uint32.valueOf(60))
+                            .setLlAdvertised(true)
+                            .setLlReceived(true)
+                            .build())
+                        .build())
+                    .build())
                 .setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol.rev151009.bgp
-                        .common.afi.safi.list.afi.safi.StateBuilder().setEnabled(false).addAugmentation(
-                                NeighborAfiSafiStateAugmentation.class, neighborAfiSafiStateAugmentation.build())
+                        .common.afi.safi.list.afi.safi.StateBuilder()
+                            .setEnabled(false)
+                            .addAugmentation(neighborAfiSafiStateAugmentation.build())
                         .build())
                 .build();
 
-        return new AfiSafisBuilder().setAfiSafi(Collections.singletonList(afiSafi)).build();
+        return new AfiSafisBuilder().setAfiSafi(Map.of(afiSafi.key(), afiSafi)).build();
     }
 
     private static ErrorHandling buildErrorHandling() {
         final ErrorHandling errorHandling = new ErrorHandlingBuilder().setState(
                 new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group.error
                         .handling.StateBuilder().setTreatAsWithdraw(false)
-                        .addAugmentation(NeighborErrorHandlingStateAugmentation.class,
-                                new NeighborErrorHandlingStateAugmentationBuilder()
+                        .addAugmentation(new NeighborErrorHandlingStateAugmentationBuilder()
                                         .setErroneousUpdateMessages(Uint32.ONE).build()).build()).build();
         return errorHandling;
     }
 
     private static Timers buildTimers() {
-        final Timers timers = new TimersBuilder().setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang
-                .bgp.rev151009.bgp.neighbor.group.timers.StateBuilder()
-                .setConnectRetry(BigDecimal.valueOf(30))
-                .setHoldTime(BigDecimal.valueOf(90))
-                .setKeepaliveInterval(BigDecimal.valueOf(30))
-                .setMinimumAdvertisementInterval(BigDecimal.valueOf(30))
-                .addAugmentation(NeighborTimersStateAugmentation.class, new NeighborTimersStateAugmentationBuilder()
-                        .setNegotiatedHoldTime(BigDecimal.TEN).setUptime(new Timeticks(Uint32.ONE)).build())
-                .build()).build();
-        return timers;
+        return new TimersBuilder()
+                .setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group
+                    .timers.StateBuilder()
+                        .setConnectRetry(BigDecimal.valueOf(30))
+                        .setHoldTime(BigDecimal.valueOf(90))
+                        .setKeepaliveInterval(BigDecimal.valueOf(30))
+                        .setMinimumAdvertisementInterval(BigDecimal.valueOf(30))
+                        .addAugmentation(new NeighborTimersStateAugmentationBuilder()
+                            .setNegotiatedHoldTime(BigDecimal.TEN)
+                            .setUptime(new Timeticks(Uint32.ONE)).build())
+                        .build())
+                .build();
     }
 
     private Transport buildTransport() {
-        final Transport transport = new TransportBuilder().setState(new org.opendaylight.yang.gen.v1.http.openconfig
+        return new TransportBuilder()
+                .setState(new org.opendaylight.yang.gen.v1.http.openconfig
                 .net.yang.bgp.rev151009.bgp.neighbor.group.transport.StateBuilder()
-                .setMtuDiscovery(false)
+                    .setMtuDiscovery(false)
                 .setPassiveMode(false)
-                .addAugmentation(NeighborTransportStateAugmentation.class,
-                        new NeighborTransportStateAugmentationBuilder().setLocalPort(this.localPort)
+                .addAugmentation(new NeighborTransportStateAugmentationBuilder()
+                    .setLocalPort(this.localPort)
                                 .setRemotePort(this.remotePort)
                                 .setRemoteAddress(new IpAddress(neighborAddress.getIpv4AddressNoZone())).build())
                 .build()).build();
-        return transport;
     }
 
     private GracefulRestart buildGracefulRestart() {
@@ -562,8 +558,7 @@ public class StateProviderImplTest extends AbstractDataBrokerTest {
         final GracefulRestart gracefulRestart = new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp
                 .rev151009.bgp.graceful.restart.GracefulRestartBuilder().setState(new org.opendaylight.yang.gen.v1.http
                 .openconfig.net.yang.bgp.rev151009.bgp.graceful.restart.graceful.restart.StateBuilder()
-                .addAugmentation(NeighborGracefulRestartStateAugmentation.class,
-                        gracefulAugmentation.build()).build()).build();
+                .addAugmentation(gracefulAugmentation.build()).build()).build();
         return gracefulRestart;
     }
 
@@ -583,11 +578,10 @@ public class StateProviderImplTest extends AbstractDataBrokerTest {
                                 .setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.multiprotocol
                                     .rev151009.bgp.common.afi.safi.list.afi.safi.StateBuilder()
                                         .setEnabled(false)
-                                        .addAugmentation(GlobalAfiSafiStateAugmentation.class,
-                                            new GlobalAfiSafiStateAugmentationBuilder()
-                                                .setTotalPaths(Uint32.valueOf(prefixesAndPaths))
-                                                .setTotalPrefixes(Uint32.valueOf(prefixesAndPaths))
-                                                .build())
+                                        .addAugmentation(new GlobalAfiSafiStateAugmentationBuilder()
+                                            .setTotalPaths(Uint32.valueOf(prefixesAndPaths))
+                                            .setTotalPrefixes(Uint32.valueOf(prefixesAndPaths))
+                                            .build())
                                         .build())
                                 .build()))
                         .build())
@@ -595,14 +589,17 @@ public class StateProviderImplTest extends AbstractDataBrokerTest {
     }
 
     private static PeerGroup buildGroupExpected() {
-        return new PeerGroupBuilder().setPeerGroupName("test-group").setState(new org.opendaylight.yang.gen.v1.http
-            .openconfig.net.yang.bgp.rev151009.bgp.neighbor.group.StateBuilder()
-            .setSendCommunity(CommunityType.NONE)
-            .setRouteFlapDamping(false)
-            .addAugmentation(PeerGroupStateAugmentation.class, new PeerGroupStateAugmentationBuilder()
-                .setTotalPaths(Uint32.ONE)
-                .setTotalPrefixes(Uint32.ONE)
-                .build()).build())
-            .build();
+        return new PeerGroupBuilder()
+                .setPeerGroupName("test-group")
+                .setState(new org.opendaylight.yang.gen.v1.http.openconfig.net.yang.bgp.rev151009.bgp.neighbor.group
+                    .StateBuilder()
+                    .setSendCommunity(CommunityType.NONE)
+                    .setRouteFlapDamping(false)
+                    .addAugmentation(new PeerGroupStateAugmentationBuilder()
+                        .setTotalPaths(Uint32.ONE)
+                        .setTotalPrefixes(Uint32.ONE)
+                        .build())
+                    .build())
+                .build();
     }
 }