BUG 5523 Add onDeviceContextLevelDown for Statistics and Rpc Managers

 *change Termination order in OFpluginProvider
 *add onDeviceContextLevelDown methods
 *add pair NodeId with Context for every manager
  *only NodeId is keyProtection with not finished unregistration
   for every shutting down order

Change-Id: Iec0defc9efeb276668404a0935114e523ab3fc10
Signed-off-by: Jozef Bacigal <jbacigal@cisco.com>

diff --git a/openflowplugin-api/src/main/java/org/opendaylight/openflowplugin/api/openflow/statistics/StatisticsContext.java b/openflowplugin-api/src/main/java/org/opendaylight/openflowplugin/api/openflow/statistics/StatisticsContext.java
index 1806b184def74568ff30e381825461f19ec4341b..7380e5656e387ff56ea8058a6022400b01781266 100644 (file)
--- a/openflowplugin-api/src/main/java/org/opendaylight/openflowplugin/api/openflow/statistics/StatisticsContext.java
+++ b/openflowplugin-api/src/main/java/org/opendaylight/openflowplugin/api/openflow/statistics/StatisticsContext.java
@@ -11,6 +11,7 @@ package org.opendaylight.openflowplugin.api.openflow.statistics;
 import com.google.common.base.Optional;
 import com.google.common.util.concurrent.ListenableFuture;
 import io.netty.util.Timeout;
+import org.opendaylight.openflowplugin.api.openflow.device.DeviceContext;
 import org.opendaylight.openflowplugin.api.openflow.device.RequestContextStack;
 import org.opendaylight.openflowplugin.api.openflow.rpc.listener.ItemLifecycleListener;
 
@@ -45,6 +46,13 @@ public interface StatisticsContext extends RequestContextStack, AutoCloseable {
      * @return dedicated item life cycle change listener (per device)
      */
     ItemLifecycleListener getItemLifeCycleListener();
+
+    /**
+     * Statistics Context has to be able to return own DeviceCtx
+     * @return {@link DeviceContext}
+     */
+    DeviceContext getDeviceContext();
+
     @Override
     void close();
 }

diff --git a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/OpenFlowPluginProviderImpl.java b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/OpenFlowPluginProviderImpl.java
index d5da30067ad6e14f44bc112171697fcfc44a4e23..d1fa7e59a04f5ad39cf297810f3ee57a0eaf7e6a 100644 (file)
--- a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/OpenFlowPluginProviderImpl.java
+++ b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/OpenFlowPluginProviderImpl.java
@@ -205,10 +205,10 @@ public class OpenFlowPluginProviderImpl implements OpenFlowPluginProvider, OpenF
         roleManager.setDeviceInitializationPhaseHandler(deviceManager);
 
         /* Termination Phase ordering - OFP Device Context suite */
-        deviceManager.setDeviceTerminationPhaseHandler(roleManager);
-        roleManager.setDeviceTerminationPhaseHandler(rpcManager);
+        deviceManager.setDeviceTerminationPhaseHandler(rpcManager);
         rpcManager.setDeviceTerminationPhaseHandler(statisticsManager);
-        statisticsManager.setDeviceTerminationPhaseHandler(deviceManager);
+        statisticsManager.setDeviceTerminationPhaseHandler(roleManager);
+        roleManager.setDeviceTerminationPhaseHandler(deviceManager);
 
         rpcManager.setStatisticsRpcEnabled(isStatisticsRpcEnabled);
         rpcManager.setNotificationPublishService(notificationPublishService);

diff --git a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/role/RoleManagerImpl.java b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/role/RoleManagerImpl.java
index 3392c9a25fee1833604ef24e28f54a0e80554dcb..ad26891fcb8b73ae2064b69df245313ffe564277 100644 (file)
--- a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/role/RoleManagerImpl.java
+++ b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/role/RoleManagerImpl.java
@@ -152,6 +152,7 @@ public class RoleManagerImpl implements RoleManager, EntityOwnershipListener {
             }
             roleContext.close();
         }
+        deviceTerminationPhaseHandler.onDeviceContextLevelDown(deviceContext);
     }
 
     private static Entity makeEntity(final NodeId nodeId) {

diff --git a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImpl.java b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImpl.java
index 027b8e9bd02e5e9a1a15ab15356dc412b9a37924..61b114056d1e4e9be5f672bc92ba39791f633f82 100644 (file)
--- a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImpl.java
+++ b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImpl.java
@@ -10,7 +10,6 @@ package org.opendaylight.openflowplugin.impl.rpc;
 import com.google.common.base.Verify;
 import com.google.common.collect.Iterators;
 import java.util.Iterator;
-import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import org.opendaylight.controller.md.sal.binding.api.NotificationPublishService;
@@ -33,7 +32,7 @@ public class RpcManagerImpl implements RpcManager {
     private DeviceInitializationPhaseHandler deviceInitPhaseHandler;
     private DeviceTerminationPhaseHandler deviceTerminPhaseHandler;
     private final int maxRequestsQuota;
-    private final ConcurrentMap<DeviceContext, RpcContext> contexts = new ConcurrentHashMap<>();
+    private final ConcurrentMap<NodeId, RpcContext> contexts = new ConcurrentHashMap<>();
     private boolean isStatisticsRpcEnabled;
     private NotificationPublishService notificationPublishService;
 
@@ -57,7 +56,7 @@ public class RpcManagerImpl implements RpcManager {
         final RpcContext rpcContext = new RpcContextImpl(deviceContext.getMessageSpy(), rpcProviderRegistry,
                 deviceContext, maxRequestsQuota, isStatisticsRpcEnabled, notificationPublishService);
 
-        Verify.verify(contexts.putIfAbsent(deviceContext, rpcContext) == null, "RPC context still not closed for node {}", nodeId);
+        Verify.verify(contexts.putIfAbsent(nodeId, rpcContext) == null, "RpcCtx still not closed for node {}", nodeId);
         deviceContext.addDeviceContextClosedHandler(this);
 
         if (OfpRole.BECOMEMASTER.equals(ofpRole)) {
@@ -80,20 +79,21 @@ public class RpcManagerImpl implements RpcManager {
 
     @Override
     public void close() {
-        for (final Iterator<Entry<DeviceContext, RpcContext>> iterator = Iterators
-                .consumingIterator(contexts.entrySet().iterator()); iterator.hasNext();) {
-            iterator.next().getValue().close();
+        for (final Iterator<RpcContext> iterator = Iterators.consumingIterator(contexts.values().iterator());
+                iterator.hasNext();) {
+            iterator.next().close();
         }
     }
 
 
     @Override
     public void onDeviceContextLevelDown(final DeviceContext deviceContext) {
-        final RpcContext removedContext = contexts.remove(deviceContext);
+        final RpcContext removedContext = contexts.remove(deviceContext.getDeviceState().getNodeId());
         if (removedContext != null) {
             LOG.info("Unregister RPCs services for device context closure");
             removedContext.close();
         }
+        deviceTerminPhaseHandler.onDeviceContextLevelDown(deviceContext);
     }
     @Override
     public void setStatisticsRpcEnabled(final boolean isStatisticsRpcEnabled) {

diff --git a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsContextImpl.java b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsContextImpl.java
index da90f9844c35a3072be31c2d7524fae2e495877e..447835bcbb1429bd8c04309968ce38712b475fa9 100644 (file)
--- a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsContextImpl.java
+++ b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsContextImpl.java
@@ -146,25 +146,26 @@ public class StatisticsContextImpl implements StatisticsContext {
         }
 
 
-        @Override
-        public <T> RequestContext<T> createRequestContext() {
-            final AbstractRequestContext<T> ret = new AbstractRequestContext<T>(deviceContext.reservedXidForDeviceMessage()) {
-                @Override
-                public void close() {
-                    requestContexts.remove(this);
-                }
-            };
-            requestContexts.add(ret);
-            return ret;
-        }
-
-        @Override
-        public void close () {
-            for (final RequestContext<?> requestContext : requestContexts) {
-                RequestContextUtil.closeRequestContextWithRpcError(requestContext, CONNECTION_CLOSED);
+    @Override
+    public <T> RequestContext<T> createRequestContext() {
+        final AbstractRequestContext<T> ret = new AbstractRequestContext<T>(deviceContext.reservedXidForDeviceMessage()) {
+            @Override
+            public void close() {
+                requestContexts.remove(this);
             }
-            if (null != pollTimeout && !pollTimeout.isExpired()) {
-                pollTimeout.cancel();
+        };
+        requestContexts.add(ret);
+        return ret;
+    }
+
+    @Override
+    public void close() {
+        for (final Iterator<RequestContext<?>> iterator = Iterators.consumingIterator(requestContexts.iterator());
+                iterator.hasNext();) {
+            RequestContextUtil.closeRequestContextWithRpcError(iterator.next(), CONNECTION_CLOSED);
+        }
+        if (null != pollTimeout && !pollTimeout.isExpired()) {
+            pollTimeout.cancel();
             }
         }
 
@@ -287,8 +288,14 @@ public class StatisticsContextImpl implements StatisticsContext {
         this.statisticsGatheringOnTheFlyService = statisticsGatheringOnTheFlyService;
     }
 
-        @Override
-        public ItemLifecycleListener getItemLifeCycleListener () {
-            return itemLifeCycleListener;
-        }
+    @Override
+    public ItemLifecycleListener getItemLifeCycleListener () {
+        return itemLifeCycleListener;
+    }
+
+
+    @Override
+    public DeviceContext getDeviceContext() {
+        return deviceContext;
     }
+}

diff --git a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImpl.java b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImpl.java
index 52c3fec570b952bf8013f2eb12c2136d247199cf..dc0661cb739d33c820cee3cd2576e3673a85bc75 100644 (file)
--- a/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImpl.java
+++ b/openflowplugin-impl/src/main/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImpl.java
@@ -29,8 +29,6 @@ import io.netty.util.HashedWheelTimer;
 import io.netty.util.Timeout;
 import io.netty.util.TimerTask;
 import java.util.Iterator;
-import java.util.Map;
-import java.util.Map.Entry;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.Future;
@@ -45,6 +43,7 @@ import org.opendaylight.openflowplugin.api.openflow.device.handlers.DeviceTermin
 import org.opendaylight.openflowplugin.api.openflow.rpc.ItemLifeCycleSource;
 import org.opendaylight.openflowplugin.api.openflow.statistics.StatisticsContext;
 import org.opendaylight.openflowplugin.api.openflow.statistics.StatisticsManager;
+import org.opendaylight.yang.gen.v1.urn.opendaylight.inventory.rev130819.NodeId;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.openflowplugin.sm.control.rev150812.ChangeStatisticsWorkModeInput;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.openflowplugin.sm.control.rev150812.GetStatisticsWorkModeOutput;
 import org.opendaylight.yang.gen.v1.urn.opendaylight.params.xml.ns.yang.openflowplugin.sm.control.rev150812.GetStatisticsWorkModeOutputBuilder;
@@ -71,7 +70,7 @@ public class StatisticsManagerImpl implements StatisticsManager, StatisticsManag
 
     private HashedWheelTimer hashedWheelTimer;
 
-    private final ConcurrentMap<DeviceContext, StatisticsContext> contexts = new ConcurrentHashMap<>();
+    private final ConcurrentMap<NodeId, StatisticsContext> contexts = new ConcurrentHashMap<>();
 
     private static final long basicTimerDelay = 3000;
     private static long currentTimerDelay = basicTimerDelay;
@@ -87,30 +86,34 @@ public class StatisticsManagerImpl implements StatisticsManager, StatisticsManag
         deviceInitPhaseHandler = handler;
     }
 
-    public StatisticsManagerImpl(@CheckForNull final RpcProviderRegistry rpcProviderRegistry, final boolean shuttingDownStatisticsPolling) {
+    public StatisticsManagerImpl(@CheckForNull final RpcProviderRegistry rpcProviderRegistry,
+                                               final boolean shuttingDownStatisticsPolling) {
         Preconditions.checkArgument(rpcProviderRegistry != null);
-        controlServiceRegistration = rpcProviderRegistry.addRpcImplementation(StatisticsManagerControlService.class, this);
+        this.controlServiceRegistration = Preconditions.checkNotNull(rpcProviderRegistry.addRpcImplementation(
+                StatisticsManagerControlService.class, this));
         this.shuttingDownStatisticsPolling = shuttingDownStatisticsPolling;
     }
 
     @Override
     public void onDeviceContextLevelUp(final DeviceContext deviceContext) throws Exception {
-        LOG.debug("Node:{}, deviceContext.getDeviceState().getRole():{}", deviceContext.getDeviceState().getNodeId(),
-                deviceContext.getDeviceState().getRole());
+        final NodeId nodeId = deviceContext.getDeviceState().getNodeId();
+        final OfpRole ofpRole = deviceContext.getDeviceState().getRole();
+        LOG.debug("Node:{}, deviceContext.getDeviceState().getRole():{}", nodeId, ofpRole);
+
         if (null == hashedWheelTimer) {
             LOG.trace("This is first device that delivered timer. Starting statistics polling immediately.");
             hashedWheelTimer = deviceContext.getTimer();
         }
         final StatisticsContext statisticsContext = new StatisticsContextImpl(deviceContext, shuttingDownStatisticsPolling);
 
-        Verify.verify(contexts.putIfAbsent(deviceContext, statisticsContext) == null, "StatisticsCtx still not closed for Node {}",deviceContext.getDeviceState().getNodeId());
+        Verify.verify(contexts.putIfAbsent(nodeId, statisticsContext) == null, "StatisticsCtx still not closed for Node {}", nodeId);
         deviceContext.addDeviceContextClosedHandler(this);
 
         if (shuttingDownStatisticsPolling) {
             LOG.info("Statistics is shutdown for node:{}", deviceContext.getDeviceState().getNodeId());
         } else {
             LOG.info("Schedule Statistics poll for node:{}", deviceContext.getDeviceState().getNodeId());
-            if (OfpRole.BECOMEMASTER.equals(deviceContext.getDeviceState().getRole())) {
+            if (OfpRole.BECOMEMASTER.equals(ofpRole)) {
                 initialStatPollForMaster(statisticsContext, deviceContext);
                 /* we want to wait for initial statCollecting response */
                 return;
@@ -136,21 +139,21 @@ public class StatisticsManagerImpl implements StatisticsManager, StatisticsManag
                         deviceInitPhaseHandler.onDeviceContextLevelUp(deviceContext);
                     } catch (final Exception e) {
                         LOG.info("failed to complete levelUp on next handler for device {}", deviceContext.getDeviceState().getNodeId());
-                        deviceContext.close();
+                        deviceContext.shutdownConnection();
                         return;
                     }
                     deviceContext.getDeviceState().setDeviceSynchronized(true);
                 } else {
                     final String deviceAddress = deviceContext.getPrimaryConnectionContext().getConnectionAdapter().getRemoteAddress().toString();
                     LOG.info("Statistics for device {} could not be gathered. Closing its device context.", deviceAddress);
-                    deviceContext.close();
+                    deviceContext.shutdownConnection();
                 }
             }
 
             @Override
             public void onFailure(final Throwable throwable) {
                 LOG.warn("Statistics manager was not able to collect dynamic info for device.", deviceContext.getDeviceState().getNodeId(), throwable);
-                deviceContext.close();
+                deviceContext.shutdownConnection();
             }
         });
     }
@@ -254,11 +257,12 @@ public class StatisticsManagerImpl implements StatisticsManager, StatisticsManag
 
     @Override
     public void onDeviceContextLevelDown(final DeviceContext deviceContext) {
-        final StatisticsContext statisticsContext = contexts.remove(deviceContext);
+        final StatisticsContext statisticsContext = contexts.remove(deviceContext.getDeviceState().getNodeId());
         if (null != statisticsContext) {
             LOG.trace("Removing device context from stack. No more statistics gathering for node {}", deviceContext.getDeviceState().getNodeId());
             statisticsContext.close();
         }
+        deviceTerminPhaseHandler.onDeviceContextLevelDown(deviceContext);
     }
 
     @Override
@@ -277,9 +281,8 @@ public class StatisticsManagerImpl implements StatisticsManager, StatisticsManag
             if (!workMode.equals(targetWorkMode)) {
                 shuttingDownStatisticsPolling = StatisticsWorkMode.FULLYDISABLED.equals(targetWorkMode);
                 // iterate through stats-ctx: propagate mode
-                for (final Map.Entry<DeviceContext, StatisticsContext> contextEntry : contexts.entrySet()) {
-                    final DeviceContext deviceContext = contextEntry.getKey();
-                    final StatisticsContext statisticsContext = contextEntry.getValue();
+                for (final StatisticsContext statisticsContext : contexts.values()) {
+                    final DeviceContext deviceContext = statisticsContext.getDeviceContext();
                     switch (targetWorkMode) {
                         case COLLECTALL:
                             scheduleNextPolling(deviceContext, statisticsContext, new TimeCounter());
@@ -318,9 +321,9 @@ public class StatisticsManagerImpl implements StatisticsManager, StatisticsManag
             controlServiceRegistration.close();
             controlServiceRegistration = null;
         }
-        for (final Iterator<Entry<DeviceContext, StatisticsContext>> iterator = Iterators
-                .consumingIterator(contexts.entrySet().iterator()); iterator.hasNext();) {
-            iterator.next().getValue().close();
+        for (final Iterator<StatisticsContext> iterator = Iterators.consumingIterator(contexts.values().iterator());
+                iterator.hasNext();) {
+            iterator.next().close();
         }
     }
 

diff --git a/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImplTest.java b/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImplTest.java
index 8f474c86e3b0712f47c42745e39bb0605ef16123..5da9d2290bdd546c85d4add083c7db4261886ede 100644 (file)
--- a/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImplTest.java
+++ b/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/rpc/RpcManagerImplTest.java
@@ -63,7 +63,8 @@ public class RpcManagerImplTest {
 
     @Before
     public void setUp() {
-        nodePath = KeyedInstanceIdentifier.create(Nodes.class).child(Node.class, new NodeKey(new NodeId("openflow-junit:1")));
+        final NodeKey nodeKey = new NodeKey(new NodeId("openflow-junit:1"));
+        nodePath = KeyedInstanceIdentifier.create(Nodes.class).child(Node.class, nodeKey);
         rpcManager = new RpcManagerImpl(rpcProviderRegistry, 5);
         rpcManager.setDeviceInitializationPhaseHandler(deviceINitializationPhaseHandler);
         FeaturesReply features = new GetFeaturesOutputBuilder()
@@ -76,6 +77,7 @@ public class RpcManagerImplTest {
         Mockito.when(deviceContext.getItemLifeCycleSourceRegistry()).thenReturn(itemLifeCycleRegistry);
         Mockito.when(deviceState.getNodeInstanceIdentifier()).thenReturn(nodePath);
         Mockito.when(deviceContext.getMessageSpy()).thenReturn(messageSpy);
+        Mockito.when(deviceState.getNodeId()).thenReturn(nodeKey.getId());
     }
 
     @Test

diff --git a/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImplTest.java b/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImplTest.java
index f6ed8e15c3e39cb4ffc624c35beab83ba6921b47..7a56c8ce9c7744e067a7d6247d4e3a65bbfd466b 100644 (file)
--- a/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImplTest.java
+++ b/openflowplugin-impl/src/test/java/org/opendaylight/openflowplugin/impl/statistics/StatisticsManagerImplTest.java
@@ -38,6 +38,7 @@ import org.opendaylight.openflowplugin.api.openflow.device.DeviceState;
 import org.opendaylight.openflowplugin.api.openflow.device.RequestContext;
 import org.opendaylight.openflowplugin.api.openflow.device.RequestContextStack;
 import org.opendaylight.openflowplugin.api.openflow.device.handlers.DeviceInitializationPhaseHandler;
+import org.opendaylight.openflowplugin.api.openflow.device.handlers.DeviceTerminationPhaseHandler;
 import org.opendaylight.openflowplugin.api.openflow.device.handlers.MultiMsgCollector;
 import org.opendaylight.openflowplugin.api.openflow.registry.ItemLifeCycleRegistry;
 import org.opendaylight.openflowplugin.api.openflow.rpc.ItemLifeCycleSource;
@@ -85,6 +86,8 @@ public class StatisticsManagerImplTest {
     @Mock
     DeviceInitializationPhaseHandler mockedDevicePhaseHandler;
     @Mock
+    DeviceTerminationPhaseHandler mockedTerminationPhaseHandler;
+    @Mock
     private RpcProviderRegistry rpcProviderRegistry;
     @Mock
     private HashedWheelTimer hashedWheelTimer;
@@ -174,7 +177,7 @@ public class StatisticsManagerImplTest {
         statisticsManager = new StatisticsManagerImpl(rpcProviderRegistry, true);
         Mockito.doAnswer(new Answer<Void>() {
             @Override
-            public Void answer(InvocationOnMock invocation) throws Throwable {
+            public Void answer(final InvocationOnMock invocation) throws Throwable {
                 final FutureCallback<OfHeader> callback = (FutureCallback<OfHeader>) invocation.getArguments()[2];
                 LOG.debug("committing entry: {}", ((MultipartRequestInput) invocation.getArguments()[1]).getType());
                 callback.onSuccess(null);
@@ -196,23 +199,26 @@ public class StatisticsManagerImplTest {
 
     @Test
     public void testOnDeviceContextClosed() throws Exception {
-        StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
-        final Map<DeviceContext, StatisticsContext> contextsMap = getContextsMap(statisticsManager);
+        final StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
+        final Map<NodeId, StatisticsContext> contextsMap = getContextsMap(statisticsManager);
 
-        contextsMap.put(mockedDeviceContext, statisticContext);
+        contextsMap.put(mockedDeviceContext.getDeviceState().getNodeId(), statisticContext);
         Assert.assertEquals(1, contextsMap.size());
 
+        statisticsManager.setDeviceTerminationPhaseHandler(mockedTerminationPhaseHandler);
         statisticsManager.onDeviceContextLevelDown(mockedDeviceContext);
         verify(statisticContext).close();
+        verify(mockedTerminationPhaseHandler).onDeviceContextLevelDown(mockedDeviceContext);
         Assert.assertEquals(0, contextsMap.size());
     }
 
-    private static Map<DeviceContext, StatisticsContext> getContextsMap(StatisticsManagerImpl statisticsManager) throws NoSuchFieldException, IllegalAccessException {
+    private static Map<NodeId, StatisticsContext> getContextsMap(final StatisticsManagerImpl statisticsManager)
+            throws NoSuchFieldException, IllegalAccessException {
         // HACK: contexts map for testing shall be accessed in some more civilized way
         final Field contextsField = StatisticsManagerImpl.class.getDeclaredField("contexts");
         Assert.assertNotNull(contextsField);
         contextsField.setAccessible(true);
-        return (Map<DeviceContext, StatisticsContext>) contextsField.get(statisticsManager);
+        return (Map<NodeId, StatisticsContext>) contextsField.get(statisticsManager);
     }
 
     @Test
@@ -231,13 +237,14 @@ public class StatisticsManagerImplTest {
      */
     @Test
     public void testChangeStatisticsWorkMode1() throws Exception {
-        StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
+        final StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
+        when(statisticContext.getDeviceContext()).thenReturn(mockedDeviceContext);
         when(statisticContext.getPollTimeout()).thenReturn(
                 Optional.<Timeout>absent());
         when(itemLifeCycleRegistry.getLifeCycleSources()).thenReturn(
                 Collections.<ItemLifeCycleSource>emptyList());
 
-        getContextsMap(statisticsManager).put(mockedDeviceContext, statisticContext);
+        getContextsMap(statisticsManager).put(mockedDeviceContext.getDeviceState().getNodeId(), statisticContext);
 
         final ChangeStatisticsWorkModeInputBuilder changeStatisticsWorkModeInputBld =
                 new ChangeStatisticsWorkModeInputBuilder()
@@ -264,15 +271,16 @@ public class StatisticsManagerImplTest {
      */
     @Test
     public void testChangeStatisticsWorkMode2() throws Exception {
-        Timeout pollTimeout = Mockito.mock(Timeout.class);
-        ItemLifeCycleSource itemLifecycleSource = Mockito.mock(ItemLifeCycleSource.class);
-        StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
+        final Timeout pollTimeout = Mockito.mock(Timeout.class);
+        final ItemLifeCycleSource itemLifecycleSource = Mockito.mock(ItemLifeCycleSource.class);
+        final StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
+        when(statisticContext.getDeviceContext()).thenReturn(mockedDeviceContext);
         when(statisticContext.getPollTimeout()).thenReturn(
                 Optional.of(pollTimeout));
         when(itemLifeCycleRegistry.getLifeCycleSources()).thenReturn(
                 Collections.singletonList(itemLifecycleSource));
 
-        getContextsMap(statisticsManager).put(mockedDeviceContext, statisticContext);
+        getContextsMap(statisticsManager).put(mockedDeviceContext.getDeviceState().getNodeId(), statisticContext);
 
         final ChangeStatisticsWorkModeInputBuilder changeStatisticsWorkModeInputBld =
                 new ChangeStatisticsWorkModeInputBuilder()
@@ -300,7 +308,8 @@ public class StatisticsManagerImplTest {
         Mockito.doNothing().when(itemLifecycleSource)
                 .setItemLifecycleListener(itemLifeCycleListenerCapt.capture());
 
-        StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
+        final StatisticsContext statisticContext = Mockito.mock(StatisticsContext.class);
+        when(statisticContext.getDeviceContext()).thenReturn(mockedDeviceContext);
         when(statisticContext.getPollTimeout()).thenReturn(
                 Optional.of(pollTimeout));
         when(statisticContext.getItemLifeCycleListener()).thenReturn(
@@ -308,7 +317,7 @@ public class StatisticsManagerImplTest {
         when(itemLifeCycleRegistry.getLifeCycleSources()).thenReturn(
                 Collections.singletonList(itemLifecycleSource));
 
-        getContextsMap(statisticsManager).put(mockedDeviceContext, statisticContext);
+        getContextsMap(statisticsManager).put(mockedDeviceContext.getDeviceState().getNodeId(), statisticContext);
 
         final ChangeStatisticsWorkModeInputBuilder changeStatisticsWorkModeInputBld =
                 new ChangeStatisticsWorkModeInputBuilder()
@@ -342,7 +351,7 @@ public class StatisticsManagerImplTest {
 
     @Test
     public void testCalculateTimerDelay() throws Exception {
-        TimeCounter timeCounter = Mockito.mock(TimeCounter.class);
+        final TimeCounter timeCounter = Mockito.mock(TimeCounter.class);
         when(timeCounter.getAverageTimeBetweenMarks()).thenReturn(2000L, 4000L);
 
         statisticsManager.calculateTimerDelay(timeCounter);