Merge java-network-refactor branch

git-svn-id: https://svn.apache.org/repos/asf/qpid/trunk/qpid@825362 13f79535-47bb-0310-9956-ffa450edef68

7 files changed, 80 insertions, 418 deletions

diff --git a/java/broker/src/test/java/org/apache/qpid/server/configuration/ServerConfigurationTest.java b/java/broker/src/test/java/org/apache/qpid/server/configuration/ServerConfigurationTest.java
index ebfa80d139..08b95f9d77 100644
--- a/java/broker/src/test/java/org/apache/qpid/server/configuration/ServerConfigurationTest.java
+++ b/java/broker/src/test/java/org/apache/qpid/server/configuration/ServerConfigurationTest.java
@@ -20,26 +20,26 @@
  */
 package org.apache.qpid.server.configuration;
 
+import java.io.File;
+import java.io.FileWriter;
+import java.io.IOException;
+import java.io.RandomAccessFile;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Locale;
+
 import junit.framework.TestCase;
+
 import org.apache.commons.configuration.Configuration;
 import org.apache.commons.configuration.ConfigurationException;
 import org.apache.commons.configuration.XMLConfiguration;
-import org.apache.qpid.codec.AMQCodecFactory;
-import org.apache.qpid.server.protocol.AMQMinaProtocolSession;
+import org.apache.qpid.server.protocol.AMQProtocolEngine;
 import org.apache.qpid.server.protocol.AMQProtocolSession;
-import org.apache.qpid.server.protocol.TestIoSession;
 import org.apache.qpid.server.registry.ApplicationRegistry;
 import org.apache.qpid.server.registry.ConfigurationFileApplicationRegistry;
 import org.apache.qpid.server.virtualhost.VirtualHost;
 import org.apache.qpid.server.virtualhost.VirtualHostRegistry;
-
-import java.io.File;
-import java.io.FileWriter;
-import java.io.IOException;
-import java.io.RandomAccessFile;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Locale;
+import org.apache.qpid.transport.TestNetworkDriver;
 
 public class ServerConfigurationTest extends TestCase
 {
@@ -589,12 +589,12 @@ public class ServerConfigurationTest extends TestCase
     {
         // Check default
         ServerConfiguration serverConfig = new ServerConfiguration(_config);
-        assertEquals(true, serverConfig.getSSLOnly());
+        assertEquals(false, serverConfig.getSSLOnly());
 
         // Check value we set
-        _config.setProperty("connector.ssl.sslOnly", false);
+        _config.setProperty("connector.ssl.sslOnly", true);
         serverConfig = new ServerConfiguration(_config);
-        assertEquals(false, serverConfig.getSSLOnly());
+        assertEquals(true, serverConfig.getSSLOnly());
     }
 
     public void testGetSSLPort() throws ConfigurationException
@@ -791,16 +791,15 @@ public class ServerConfigurationTest extends TestCase
         // Test config
         VirtualHostRegistry virtualHostRegistry = reg.getVirtualHostRegistry();
         VirtualHost virtualHost = virtualHostRegistry.getVirtualHost("test");
-        AMQCodecFactory codecFactory = new AMQCodecFactory(true);
 
-        TestIoSession iosession = new TestIoSession();
-        iosession.setAddress("127.0.0.1");
+        TestNetworkDriver testDriver = new TestNetworkDriver();
+        testDriver.setRemoteAddress("127.0.0.1");
 
-        AMQProtocolSession session = new AMQMinaProtocolSession(iosession, virtualHostRegistry, codecFactory);
+        AMQProtocolEngine session = new AMQProtocolEngine(virtualHostRegistry, testDriver);
         assertFalse(reg.getAccessManager().authoriseConnect(session, virtualHost));
 
-        iosession.setAddress("127.1.2.3");
-        session = new AMQMinaProtocolSession(iosession, virtualHostRegistry, codecFactory);
+        testDriver.setRemoteAddress("127.1.2.3");
+        session = new AMQProtocolEngine(virtualHostRegistry, testDriver);
         assertTrue(reg.getAccessManager().authoriseConnect(session, virtualHost));
     }
 
@@ -866,12 +865,12 @@ public class ServerConfigurationTest extends TestCase
         // Test config
         VirtualHostRegistry virtualHostRegistry = reg.getVirtualHostRegistry();
         VirtualHost virtualHost = virtualHostRegistry.getVirtualHost("test");
-        AMQCodecFactory codecFactory = new AMQCodecFactory(true);
 
-        TestIoSession iosession = new TestIoSession();
-        iosession.setAddress("127.0.0.1");
+        TestNetworkDriver testDriver = new TestNetworkDriver();
+        testDriver.setRemoteAddress("127.0.0.1");
 
-        AMQProtocolSession session = new AMQMinaProtocolSession(iosession, virtualHostRegistry, codecFactory);
+        AMQProtocolEngine session = new AMQProtocolEngine(virtualHostRegistry, testDriver);
+        session.setNetworkDriver(testDriver);
         assertFalse(reg.getAccessManager().authoriseConnect(session, virtualHost));
     }
 
@@ -935,12 +934,11 @@ public class ServerConfigurationTest extends TestCase
         ApplicationRegistry.initialise(reg, 1);
 
         // Test config
-        TestIoSession iosession = new TestIoSession();
-        iosession.setAddress("127.0.0.1");
+        TestNetworkDriver testDriver = new TestNetworkDriver();
+        testDriver.setRemoteAddress("127.0.0.1");
         VirtualHostRegistry virtualHostRegistry = reg.getVirtualHostRegistry();
         VirtualHost virtualHost = virtualHostRegistry.getVirtualHost("test");
-        AMQCodecFactory codecFactory = new AMQCodecFactory(true);
-        AMQProtocolSession session = new AMQMinaProtocolSession(iosession, virtualHostRegistry, codecFactory);
+        AMQProtocolSession session = new AMQProtocolEngine(virtualHostRegistry, testDriver);
         assertFalse(reg.getAccessManager().authoriseConnect(session, virtualHost));
 
         RandomAccessFile fileBRandom = new RandomAccessFile(fileB, "rw");
diff --git a/java/broker/src/test/java/org/apache/qpid/server/protocol/AMQProtocolSessionMBeanTest.java b/java/broker/src/test/java/org/apache/qpid/server/protocol/AMQProtocolSessionMBeanTest.java
index e199255f50..bc36c61382 100644
--- a/java/broker/src/test/java/org/apache/qpid/server/protocol/AMQProtocolSessionMBeanTest.java
+++ b/java/broker/src/test/java/org/apache/qpid/server/protocol/AMQProtocolSessionMBeanTest.java
@@ -44,7 +44,7 @@ public class AMQProtocolSessionMBeanTest extends TestCase
     private static final Logger log = Logger.getLogger(AMQProtocolSessionMBeanTest.class);
 
     private MessageStore _messageStore = new SkeletonMessageStore();
-    private AMQMinaProtocolSession _protocolSession;
+    private AMQProtocolEngine _protocolSession;
     private AMQChannel _channel;
     private AMQProtocolSessionMBean _mbean;
 
diff --git a/java/broker/src/test/java/org/apache/qpid/server/protocol/InternalTestProtocolSession.java b/java/broker/src/test/java/org/apache/qpid/server/protocol/InternalTestProtocolSession.java
index 37dfead2e5..ec7bf1cb72 100644
--- a/java/broker/src/test/java/org/apache/qpid/server/protocol/InternalTestProtocolSession.java
+++ b/java/broker/src/test/java/org/apache/qpid/server/protocol/InternalTestProtocolSession.java
@@ -20,23 +20,23 @@
  */
 package org.apache.qpid.server.protocol;
 
-import org.apache.qpid.AMQException;
-import org.apache.qpid.codec.AMQCodecFactory;
-import org.apache.qpid.framing.AMQShortString;
-import org.apache.qpid.server.output.ProtocolOutputConverter;
-import org.apache.qpid.server.queue.AMQMessage;
-import org.apache.qpid.server.registry.ApplicationRegistry;
-import org.apache.qpid.server.virtualhost.VirtualHost;
-
+import java.security.Principal;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.security.Principal;
 
-public class InternalTestProtocolSession extends AMQMinaProtocolSession implements ProtocolOutputConverter
+import org.apache.qpid.AMQException;
+import org.apache.qpid.framing.AMQShortString;
+import org.apache.qpid.server.output.ProtocolOutputConverter;
+import org.apache.qpid.server.queue.AMQMessage;
+import org.apache.qpid.server.registry.ApplicationRegistry;
+import org.apache.qpid.server.virtualhost.VirtualHost;
+import org.apache.qpid.transport.TestNetworkDriver;
+
+public class InternalTestProtocolSession extends AMQProtocolEngine implements ProtocolOutputConverter
 {
     // ChannelID(LIST)  -> LinkedList<Pair>
     final Map<Integer, Map<AMQShortString, LinkedList<DeliveryPair>>> _channelDelivers;
@@ -44,9 +44,7 @@ public class InternalTestProtocolSession extends AMQMinaProtocolSession implemen
 
     public InternalTestProtocolSession(VirtualHost virtualHost) throws AMQException
     {
-        super(new TestIoSession(),
-              ApplicationRegistry.getInstance().getVirtualHostRegistry(),
-              new AMQCodecFactory(true));
+        super(ApplicationRegistry.getInstance().getVirtualHostRegistry(), new TestNetworkDriver());
 
         _channelDelivers = new HashMap<Integer, Map<AMQShortString, LinkedList<DeliveryPair>>>();
 
diff --git a/java/broker/src/test/java/org/apache/qpid/server/protocol/MaxChannelsTest.java b/java/broker/src/test/java/org/apache/qpid/server/protocol/MaxChannelsTest.java
index 8597fc5863..e37492bcb0 100644
--- a/java/broker/src/test/java/org/apache/qpid/server/protocol/MaxChannelsTest.java
+++ b/java/broker/src/test/java/org/apache/qpid/server/protocol/MaxChannelsTest.java
@@ -37,7 +37,7 @@ import java.security.Principal;
 /** Test class to test MBean operations for AMQMinaProtocolSession. */
 public class MaxChannelsTest extends TestCase
 {
-	private AMQMinaProtocolSession _session;
+	private AMQProtocolEngine _session;
 
     public void testChannels() throws Exception
     {
diff --git a/java/broker/src/test/java/org/apache/qpid/server/protocol/TestIoSession.java b/java/broker/src/test/java/org/apache/qpid/server/protocol/TestIoSession.java
deleted file mode 100644
index 211f491867..0000000000
--- a/java/broker/src/test/java/org/apache/qpid/server/protocol/TestIoSession.java
+++ /dev/null
@@ -1,328 +0,0 @@
-/*
- *
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- *
- */
-package org.apache.qpid.server.protocol;
-
-import java.net.InetSocketAddress;
-import java.net.SocketAddress;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-
-import org.apache.mina.common.CloseFuture;
-import org.apache.mina.common.IdleStatus;
-import org.apache.mina.common.IoFilterChain;
-import org.apache.mina.common.IoHandler;
-import org.apache.mina.common.IoService;
-import org.apache.mina.common.IoServiceConfig;
-import org.apache.mina.common.IoSession;
-import org.apache.mina.common.IoSessionConfig;
-import org.apache.mina.common.ThreadModel;
-import org.apache.mina.common.TrafficMask;
-import org.apache.mina.common.TransportType;
-import org.apache.mina.common.WriteFuture;
-import org.apache.mina.transport.socket.nio.SocketAcceptorConfig;
-import org.apache.qpid.pool.ReadWriteThreadModel;
-
-/**
- * Test implementation of IoSession, which is required for some tests. Methods not being used are not implemented,
- * so if this class is being used and some methods are to be used, then please update those.
- */
-public class TestIoSession implements IoSession
-{
-    private final ConcurrentMap attributes = new ConcurrentHashMap();
-    private String _address = "127.0.0.1";
-    private int _port = 1;
-
-    public TestIoSession()
-    {
-    }
-
-    public IoService getService()
-    {
-        return null;
-    }
-
-    public IoServiceConfig getServiceConfig()
-    {
-        return new TestIoConfig();
-    }
-
-    public IoHandler getHandler()
-    {
-        return null;
-    }
-
-    public IoSessionConfig getConfig()
-    {
-        return null;
-    }
-
-    public IoFilterChain getFilterChain()
-    {
-        return null;
-    }
-
-    public WriteFuture write(Object message)
-    {
-        return null;
-    }
-
-    public CloseFuture close()
-    {
-        return null;
-    }
-
-    public Object getAttachment()
-    {
-        return getAttribute("");
-    }
-
-    public Object setAttachment(Object attachment)
-    {
-        return setAttribute("",attachment);
-    }
-
-    public Object getAttribute(String key)
-    {
-        return attributes.get(key);
-    }
-
-    public Object setAttribute(String key, Object value)
-    {
-        return attributes.put(key,value);
-    }
-
-    public Object setAttribute(String key)
-    {
-        return attributes.put(key, Boolean.TRUE);
-    }
-
-    public Object removeAttribute(String key)
-    {
-        return attributes.remove(key);
-    }
-
-    public boolean containsAttribute(String key)
-    {
-        return attributes.containsKey(key);
-    }
-
-    public Set getAttributeKeys()
-    {
-        return attributes.keySet();
-    }
-
-    public TransportType getTransportType()
-    {
-        return null;
-    }
-
-    public boolean isConnected()
-    {
-        return false;
-    }
-
-    public boolean isClosing()
-    {
-        return false;
-    }
-
-    public CloseFuture getCloseFuture()
-    {
-        return null;
-    }
-
-    public SocketAddress getRemoteAddress()
-    {
-        return new InetSocketAddress(getAddress(), getPort());
-    }
-
-    public SocketAddress getLocalAddress()
-    {
-        return null;
-    }
-
-    public SocketAddress getServiceAddress()
-    {
-        return null;
-    }
-
-    public int getIdleTime(IdleStatus status)
-    {
-        return 0;
-    }
-
-    public long getIdleTimeInMillis(IdleStatus status)
-    {
-        return 0;
-    }
-
-    public void setIdleTime(IdleStatus status, int idleTime)
-    {
-
-    }
-
-    public int getWriteTimeout()
-    {
-        return 0;
-    }
-
-    public long getWriteTimeoutInMillis()
-    {
-        return 0;
-    }
-
-    public void setWriteTimeout(int writeTimeout)
-    {
-
-    }
-
-    public TrafficMask getTrafficMask()
-    {
-        return null;
-    }
-
-    public void setTrafficMask(TrafficMask trafficMask)
-    {
-
-    }
-
-    public void suspendRead()
-    {
-
-    }
-
-    public void suspendWrite()
-    {
-
-    }
-
-    public void resumeRead()
-    {
-
-    }
-
-    public void resumeWrite()
-    {
-
-    }
-
-    public long getReadBytes()
-    {
-        return 0;
-    }
-
-    public long getWrittenBytes()
-    {
-        return 0;
-    }
-
-    public long getReadMessages()
-    {
-        return 0;
-    }
-
-    public long getWrittenMessages()
-    {
-        return 0;
-    }
-
-    public long getWrittenWriteRequests()
-    {
-        return 0;
-    }
-
-    public int getScheduledWriteRequests()
-    {
-        return 0;
-    }
-
-    public int getScheduledWriteBytes()
-    {
-        return 0;
-    }
-
-    public long getCreationTime()
-    {
-        return 0;
-    }
-
-    public long getLastIoTime()
-    {
-        return 0;
-    }
-
-    public long getLastReadTime()
-    {
-        return 0;
-    }
-
-    public long getLastWriteTime()
-    {
-        return 0;
-    }
-
-    public boolean isIdle(IdleStatus status)
-    {
-        return false;
-    }
-
-    public int getIdleCount(IdleStatus status)
-    {
-        return 0;
-    }
-
-    public long getLastIdleTime(IdleStatus status)
-    {
-        return 0; 
-    }
-
-    public void setAddress(String string)
-    {
-        this._address = string;
-    }
-
-    public String getAddress()
-    {
-        return _address;
-    }
-
-    public void setPort(int _port)
-    {
-        this._port = _port;
-    }
-
-    public int getPort()
-    {
-        return _port;
-    }
-
-    /**
-     * Test implementation of IoServiceConfig
-     */
-    private class TestIoConfig extends SocketAcceptorConfig
-    {
-        public ThreadModel getThreadModel()
-        {
-            return ReadWriteThreadModel.getInstance();
-        }
-    }
-}
diff --git a/java/broker/src/test/java/org/apache/qpid/server/queue/AMQQueueAlertTest.java b/java/broker/src/test/java/org/apache/qpid/server/queue/AMQQueueAlertTest.java
index 8c6260ca9e..19470e6226 100644
--- a/java/broker/src/test/java/org/apache/qpid/server/queue/AMQQueueAlertTest.java
+++ b/java/broker/src/test/java/org/apache/qpid/server/queue/AMQQueueAlertTest.java
@@ -20,37 +20,34 @@
  */
 package org.apache.qpid.server.queue;
 
+import java.util.ArrayList;
+import java.util.LinkedList;
+
+import javax.management.Notification;
+
 import junit.framework.TestCase;
+
+import org.apache.mina.common.ByteBuffer;
 import org.apache.qpid.AMQException;
-import org.apache.qpid.server.store.MessageStore;
-import org.apache.qpid.server.store.MemoryMessageStore;
-import org.apache.qpid.server.store.StoreContext;
-import org.apache.qpid.server.virtualhost.VirtualHost;
-import org.apache.qpid.server.registry.IApplicationRegistry;
-import org.apache.qpid.server.registry.ApplicationRegistry;
-import org.apache.qpid.server.txn.TransactionalContext;
-import org.apache.qpid.server.txn.NonTransactionalContext;
-import org.apache.qpid.server.RequiredDeliveryException;
+import org.apache.qpid.framing.AMQShortString;
+import org.apache.qpid.framing.ContentHeaderBody;
+import org.apache.qpid.framing.abstraction.ContentChunk;
+import org.apache.qpid.framing.abstraction.MessagePublishInfo;
 import org.apache.qpid.server.AMQChannel;
+import org.apache.qpid.server.RequiredDeliveryException;
 import org.apache.qpid.server.logging.actors.CurrentActor;
+import org.apache.qpid.server.protocol.AMQProtocolEngine;
+import org.apache.qpid.server.protocol.InternalTestProtocolSession;
+import org.apache.qpid.server.registry.ApplicationRegistry;
+import org.apache.qpid.server.registry.IApplicationRegistry;
+import org.apache.qpid.server.store.MemoryMessageStore;
+import org.apache.qpid.server.store.MessageStore;
+import org.apache.qpid.server.store.StoreContext;
 import org.apache.qpid.server.subscription.Subscription;
 import org.apache.qpid.server.subscription.SubscriptionFactoryImpl;
-import org.apache.qpid.server.protocol.AMQMinaProtocolSession;
-import org.apache.qpid.server.protocol.InternalTestProtocolSession;
-import org.apache.qpid.framing.ContentHeaderBody;
-import org.apache.qpid.framing.AMQShortString;
-import org.apache.qpid.framing.abstraction.MessagePublishInfo;
-import org.apache.qpid.framing.abstraction.ContentChunk;
-import org.apache.commons.configuration.CompositeConfiguration;
-import org.apache.mina.common.ByteBuffer;
-
-import javax.management.Notification;
-
-import java.util.ArrayList;
-import java.util.LinkedList;
-import java.util.Collections;
-import java.util.Set;
-import java.security.Principal;
+import org.apache.qpid.server.txn.NonTransactionalContext;
+import org.apache.qpid.server.txn.TransactionalContext;
+import org.apache.qpid.server.virtualhost.VirtualHost;
 
 /** This class tests all the alerts an AMQQueue can throw based on threshold values of different parameters */
 public class AMQQueueAlertTest extends TestCase
@@ -62,7 +59,7 @@ public class AMQQueueAlertTest extends TestCase
     private AMQQueue _queue;
     private AMQQueueMBean _queueMBean;
     private VirtualHost _virtualHost;
-    private AMQMinaProtocolSession _protocolSession;
+    private AMQProtocolEngine _protocolSession;
     private MessageStore _messageStore = new MemoryMessageStore();
     private StoreContext _storeContext = new StoreContext();
     private TransactionalContext _transactionalContext = new NonTransactionalContext(_messageStore, _storeContext,
diff --git a/java/broker/src/test/java/org/apache/qpid/server/security/access/plugins/network/FirewallPluginTest.java b/java/broker/src/test/java/org/apache/qpid/server/security/access/plugins/network/FirewallPluginTest.java
index a497365b06..5d3335c001 100644
--- a/java/broker/src/test/java/org/apache/qpid/server/security/access/plugins/network/FirewallPluginTest.java
+++ b/java/broker/src/test/java/org/apache/qpid/server/security/access/plugins/network/FirewallPluginTest.java
@@ -30,17 +30,14 @@ import java.net.InetSocketAddress;
 import junit.framework.TestCase;
 
 import org.apache.commons.configuration.ConfigurationException;
-import org.apache.commons.configuration.PropertiesConfiguration;
 import org.apache.commons.configuration.XMLConfiguration;
-import org.apache.qpid.codec.AMQCodecFactory;
-import org.apache.qpid.server.configuration.VirtualHostConfiguration;
-import org.apache.qpid.server.protocol.AMQMinaProtocolSession;
-import org.apache.qpid.server.protocol.TestIoSession;
+import org.apache.qpid.server.protocol.AMQProtocolEngine;
+import org.apache.qpid.server.registry.ApplicationRegistry;
 import org.apache.qpid.server.security.access.ACLPlugin.AuthzResult;
 import org.apache.qpid.server.store.TestableMemoryMessageStore;
 import org.apache.qpid.server.virtualhost.VirtualHost;
 import org.apache.qpid.server.virtualhost.VirtualHostRegistry;
-import org.apache.qpid.server.registry.ApplicationRegistry;
+import org.apache.qpid.transport.TestNetworkDriver;
 
 public class FirewallPluginTest extends TestCase
 {
@@ -84,22 +81,22 @@ public class FirewallPluginTest extends TestCase
 
     private TestableMemoryMessageStore _store;
     private VirtualHost _virtualHost;
-    private AMQMinaProtocolSession _session;
+    private AMQProtocolEngine _session;
+    private TestNetworkDriver _testDriver;
 
     @Override
     public void setUp() throws Exception
     {
         super.setUp();
         _store = new TestableMemoryMessageStore();
-        TestIoSession iosession = new TestIoSession();
-        iosession.setAddress("127.0.0.1");
+        _testDriver = new TestNetworkDriver();
+        _testDriver.setRemoteAddress("127.0.0.1");
 
         // Retreive VirtualHost from the Registry
         VirtualHostRegistry virtualHostRegistry = ApplicationRegistry.getInstance().getVirtualHostRegistry();
         _virtualHost = virtualHostRegistry.getVirtualHost("test");
 
-        AMQCodecFactory codecFactory = new AMQCodecFactory(true);
-        _session = new AMQMinaProtocolSession(iosession, virtualHostRegistry, codecFactory);        
+        _session = new AMQProtocolEngine(virtualHostRegistry, _testDriver);
     }
 
     public void tearDown() throws Exception
@@ -170,7 +167,7 @@ public class FirewallPluginTest extends TestCase
         assertEquals(AuthzResult.DENIED, plugin.authoriseConnect(_session, _virtualHost));
         
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("192.168.23.23");
+        _testDriver.setRemoteAddress("192.168.23.23");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
     
@@ -185,7 +182,7 @@ public class FirewallPluginTest extends TestCase
         assertEquals(AuthzResult.DENIED, plugin.authoriseConnect(_session, _virtualHost));
         
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("192.168.23.23");
+        _testDriver.setRemoteAddress("192.168.23.23");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
 
@@ -198,7 +195,7 @@ public class FirewallPluginTest extends TestCase
         FirewallPlugin plugin = initialisePlugin("deny", new RuleInfo[]{rule});
 
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("127.0.0.1");
+        _testDriver.setRemoteAddress("127.0.0.1");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
 
@@ -211,7 +208,7 @@ public class FirewallPluginTest extends TestCase
         FirewallPlugin plugin = initialisePlugin("deny", new RuleInfo[]{rule});
 
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("127.0.0.1");
+        _testDriver.setRemoteAddress("127.0.0.1");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
     
@@ -234,7 +231,7 @@ public class FirewallPluginTest extends TestCase
         assertEquals(AuthzResult.DENIED, plugin.authoriseConnect(_session, _virtualHost));
         
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("192.168.23.23");
+        _testDriver.setRemoteAddress("192.168.23.23");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
     
@@ -257,7 +254,7 @@ public class FirewallPluginTest extends TestCase
         assertEquals(AuthzResult.DENIED, plugin.authoriseConnect(_session, _virtualHost));
         
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("192.168.23.23");
+        _testDriver.setRemoteAddress("192.168.23.23");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
 
@@ -271,7 +268,7 @@ public class FirewallPluginTest extends TestCase
         assertEquals(AuthzResult.DENIED, plugin.authoriseConnect(_session, _virtualHost));
         
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("192.168.23.23");
+        _testDriver.setRemoteAddress("192.168.23.23");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
     
@@ -285,7 +282,7 @@ public class FirewallPluginTest extends TestCase
         assertEquals(AuthzResult.DENIED, plugin.authoriseConnect(_session, _virtualHost));
         
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("192.168.23.23");
+        _testDriver.setRemoteAddress("192.168.23.23");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }
     
@@ -295,11 +292,11 @@ public class FirewallPluginTest extends TestCase
         firstRule.setAccess("allow");
         firstRule.setHostname("foo, bar, "+new InetSocketAddress("127.0.0.1", 5672).getHostName());
         FirewallPlugin plugin = initialisePlugin("deny", new RuleInfo[]{firstRule});
-        ((TestIoSession) _session.getIOSession()).setAddress("10.0.0.1");
+        _testDriver.setRemoteAddress("10.0.0.1");
         assertEquals(AuthzResult.DENIED, plugin.authoriseConnect(_session, _virtualHost));
         
         // Set session IP so that we're connected from the right address
-        ((TestIoSession) _session.getIOSession()).setAddress("127.0.0.1");
+        _testDriver.setRemoteAddress("127.0.0.1");
         assertEquals(AuthzResult.ALLOWED, plugin.authoriseConnect(_session, _virtualHost));
     }