diff --git a/src/test/java/net/floodlightcontroller/core/internal/OFChannelHandlerTest.java b/src/test/java/net/floodlightcontroller/core/internal/OFChannelHandlerTest.java
index 173ee9f7cdc1a7db53679f2509d94d900faf570f..818c08eab57549335452af249a6d113b92c76443 100644
--- a/src/test/java/net/floodlightcontroller/core/internal/OFChannelHandlerTest.java
+++ b/src/test/java/net/floodlightcontroller/core/internal/OFChannelHandlerTest.java
@@ -696,6 +696,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+               .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, true);
         expectLastCall().once();
         sw.setHARole(Role.MASTER);
@@ -739,6 +741,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, true);
         expectLastCall().once();
         sw.setHARole(Role.SLAVE);
@@ -783,6 +787,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, false);
         expectLastCall().once();
         sw.setHARole(Role.MASTER);
@@ -844,6 +850,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, false);
         expectLastCall().once();
         sw.setHARole(Role.MASTER);
@@ -893,6 +901,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, false);
         expectLastCall().once();
         sw.setHARole(Role.SLAVE);
@@ -950,6 +960,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, false);
         expectLastCall().once();
         sw.setHARole(Role.SLAVE);
@@ -1005,6 +1017,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, true);
         expectLastCall().once();
         sw.setHARole(Role.MASTER);
@@ -1038,6 +1052,8 @@ public class OFChannelHandlerTest {
 
         // prepare mocks and inject the role reply message
         reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
         sw.setAttribute(IOFSwitch.SWITCH_SUPPORTS_NX_ROLE, true);
         expectLastCall().once();
         sw.setHARole(Role.SLAVE);
@@ -1103,6 +1119,11 @@ public class OFChannelHandlerTest {
         OFMessage err = getErrorMessage(OFErrorType.OFPET_BAD_ACTION,
                                         0,
                                         xid);
+        verify(sw);
+        reset(sw);
+        expect(sw.inputThrottled(anyObject(OFMessage.class)))
+                .andReturn(false).anyTimes();
+        replay(sw);
         sendMessageToHandlerWithControllerReset(Collections.singletonList(err));
 
         verifyExceptionCaptured(SwitchStateException.class);
diff --git a/src/test/java/net/floodlightcontroller/core/internal/OFSwitchBaseTest.java b/src/test/java/net/floodlightcontroller/core/internal/OFSwitchBaseTest.java
index 4ef843c7a170d8b2f5ee9238551bcc5e9a713277..fe71eee3ccb2f7aad70e844398e82812dfb9aff6 100644
--- a/src/test/java/net/floodlightcontroller/core/internal/OFSwitchBaseTest.java
+++ b/src/test/java/net/floodlightcontroller/core/internal/OFSwitchBaseTest.java
@@ -84,6 +84,7 @@ public class OFSwitchBaseTest {
             return null;
         }
 
+        @Override
         public void write(OFMessage msg, FloodlightContext cntx) {
             blockMessage = msg;
         }
@@ -134,7 +135,7 @@ public class OFSwitchBaseTest {
         sw = new OFSwitchTest(floodlightProvider);
         switches = new ConcurrentHashMap<Long, IOFSwitch>();
         switches.put(sw.getId(), sw);
-        expect(floodlightProvider.getSwitches()).andReturn(switches).anyTimes();
+        expect(floodlightProvider.getSwitch(sw.getId())).andReturn(sw).anyTimes();
         expect(floodlightProvider.getOFMessageFactory())
                 .andReturn(BasicFactory.getInstance()).anyTimes();
         replay(floodlightProvider);