1 files changed, 10 insertions, 6 deletions
diff --git a/java/broker/src/test/java/org/apache/qpid/server/exchange/HeadersExchangeTest.java b/java/broker/src/test/java/org/apache/qpid/server/exchange/HeadersExchangeTest.java
index 326d36df05..bd6a02d69b 100644
--- a/java/broker/src/test/java/org/apache/qpid/server/exchange/HeadersExchangeTest.java
+++ b/java/broker/src/test/java/org/apache/qpid/server/exchange/HeadersExchangeTest.java
@@ -23,8 +23,7 @@ package org.apache.qpid.server.exchange;
 import org.apache.qpid.AMQException;
 import org.apache.qpid.server.protocol.AMQProtocolSession;
 import org.apache.qpid.server.protocol.InternalTestProtocolSession;
-import org.apache.qpid.server.registry.ApplicationRegistry;
-import org.apache.qpid.server.virtualhost.VirtualHost;
+import org.apache.qpid.server.util.BrokerTestHelper;
 
 public class HeadersExchangeTest extends AbstractHeadersExchangeTestBase
 {
@@ -34,10 +33,15 @@ public class HeadersExchangeTest extends AbstractHeadersExchangeTestBase
     public void setUp() throws Exception
     {
         super.setUp();
-        // Just use the first vhost.
-        VirtualHost
-                virtualHost = ApplicationRegistry.getInstance().getVirtualHostRegistry().getVirtualHosts().iterator().next();
-        _protocolSession = new InternalTestProtocolSession(virtualHost);
+        BrokerTestHelper.setUp();
+        _protocolSession = new InternalTestProtocolSession(getVirtualHost(), BrokerTestHelper.createBrokerMock());
+    }
+
+    @Override
+    public void tearDown() throws Exception
+    {
+        BrokerTestHelper.tearDown();
+        super.tearDown();
     }
 
     public void testSimple() throws AMQException