Refactor TestWithPeerGroup a bit to make it more reusable in apps.

core/src/main/java/com/google/bitcoin/testing/TestWithPeerGroup.java

@@ -14,20 +14,20 @@
  * limitations under the License.
  */
 
-package com.google.bitcoin.core;
+package com.google.bitcoin.testing;
 
+import com.google.bitcoin.core.*;
 import com.google.bitcoin.net.BlockingClientManager;
 import com.google.bitcoin.net.NioClientManager;
 import com.google.bitcoin.params.UnitTestParams;
 import com.google.bitcoin.store.BlockStore;
-import com.google.bitcoin.testing.InboundMessageQueuer;
+import com.google.bitcoin.store.MemoryBlockStore;
-import com.google.bitcoin.testing.TestWithNetworkConnections;
 import com.google.common.base.Preconditions;
 
 import java.net.InetSocketAddress;
 
 import static com.google.common.base.Preconditions.checkArgument;
-import static org.junit.Assert.assertTrue;
+import static com.google.common.base.Preconditions.checkState;
 
 /**
  * Utility class that makes it easy to work with mock NetworkConnections in PeerGroups.
@@ -46,6 +46,11 @@ public class TestWithPeerGroup extends TestWithNetworkConnections {
         this.clientType = clientType;
     }
 
+    @Override
+    public void setUp() throws Exception {
+        setUp(new MemoryBlockStore(params));
+    }
+
     public void setUp(BlockStore blockStore) throws Exception {
         super.setUp(blockStore);
 
@@ -55,12 +60,25 @@ public class TestWithPeerGroup extends TestWithNetworkConnections {
         initPeerGroup();
     }
 
+    @Override
+    public void tearDown() {
+        try {
+            super.tearDown();
+            Utils.finishMockSleep();
+            peerGroup.stopAsync();
+            peerGroup.awaitTerminated();
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     protected void initPeerGroup() {
         if (clientType == ClientType.NIO_CLIENT_MANAGER)
             peerGroup = new PeerGroup(unitTestParams, blockChain, new NioClientManager());
         else
             peerGroup = new PeerGroup(unitTestParams, blockChain, new BlockingClientManager());
         peerGroup.setPingIntervalMsec(0);  // Disable the pings as they just get in the way of most tests.
+        peerGroup.addWallet(wallet);
     }
 
     protected InboundMessageQueuer connectPeerWithoutVersionExchange(int id) throws Exception {
@@ -82,12 +100,7 @@ public class TestWithPeerGroup extends TestWithNetworkConnections {
         // Complete handshake with the peer - send/receive version(ack)s, receive bloom filter
         writeTarget.sendMessage(versionMessage);
         writeTarget.sendMessage(new VersionAck());
-        assertTrue(writeTarget.nextMessageBlocking() instanceof VersionMessage);
+        stepThroughInit(versionMessage, writeTarget);
-        assertTrue(writeTarget.nextMessageBlocking() instanceof VersionAck);
-        if (versionMessage.isBloomFilteringSupported()) {
-            assertTrue(writeTarget.nextMessageBlocking() instanceof BloomFilter);
-            assertTrue(writeTarget.nextMessageBlocking() instanceof MemoryPoolMessage);
-        }
         return writeTarget;
     }
 
@@ -103,12 +116,16 @@ public class TestWithPeerGroup extends TestWithNetworkConnections {
         // Complete handshake with the peer - send/receive version(ack)s, receive bloom filter
         writeTarget.sendMessage(versionMessage);
         writeTarget.sendMessage(new VersionAck());
-        assertTrue(writeTarget.nextMessageBlocking() instanceof VersionMessage);
+        stepThroughInit(versionMessage, writeTarget);
-        assertTrue(writeTarget.nextMessageBlocking() instanceof VersionAck);
-        if (versionMessage.isBloomFilteringSupported()) {
-            assertTrue(writeTarget.nextMessageBlocking() instanceof BloomFilter);
-            assertTrue(writeTarget.nextMessageBlocking() instanceof MemoryPoolMessage);
-        }
         return writeTarget;
     }
+
+    private void stepThroughInit(VersionMessage versionMessage, InboundMessageQueuer writeTarget) throws InterruptedException {
+        checkState(writeTarget.nextMessageBlocking() instanceof VersionMessage);
+        checkState(writeTarget.nextMessageBlocking() instanceof VersionAck);
+        if (versionMessage.isBloomFilteringSupported()) {
+            checkState(writeTarget.nextMessageBlocking() instanceof BloomFilter);
+            checkState(writeTarget.nextMessageBlocking() instanceof MemoryPoolMessage);
+        }
+    }
 }

core/src/test/java/com/google/bitcoin/core/FilteredBlockAndPartialMerkleTreeTests.java

@@ -21,6 +21,7 @@ import com.google.bitcoin.core.TransactionConfidence.ConfidenceType;
 import com.google.bitcoin.params.UnitTestParams;
 import com.google.bitcoin.store.MemoryBlockStore;
 import com.google.bitcoin.testing.InboundMessageQueuer;
+import com.google.bitcoin.testing.TestWithPeerGroup;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;

core/src/test/java/com/google/bitcoin/core/PeerGroupTest.java

@@ -20,9 +20,9 @@ package com.google.bitcoin.core;
 import com.google.bitcoin.net.discovery.PeerDiscovery;
 import com.google.bitcoin.net.discovery.PeerDiscoveryException;
 import com.google.bitcoin.params.UnitTestParams;
-import com.google.bitcoin.store.MemoryBlockStore;
 import com.google.bitcoin.testing.FakeTxBuilder;
 import com.google.bitcoin.testing.InboundMessageQueuer;
+import com.google.bitcoin.testing.TestWithPeerGroup;
 import com.google.bitcoin.utils.Threading;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
@@ -70,6 +70,7 @@ public class PeerGroupTest extends TestWithPeerGroup {
     @Override
     @Before
     public void setUp() throws Exception {
+        super.setUp();
         peerToMessageCount = new HashMap<Peer, AtomicInteger>();
         connectedPeers = new LinkedBlockingQueue<Peer>();
         disconnectedPeers = new LinkedBlockingQueue<Peer>();
@@ -96,16 +97,12 @@ public class PeerGroupTest extends TestWithPeerGroup {
                 return m;
             }
         };
-        super.setUp(new MemoryBlockStore(UnitTestParams.get()));
-        peerGroup.addWallet(wallet);
     }
 
+    @Override
     @After
-    public void tearDown() throws Exception {
+    public void tearDown() {
         super.tearDown();
-        Utils.finishMockSleep();
-        peerGroup.stopAsync();
-        peerGroup.awaitTerminated();
     }
 
     @Test

core/src/test/java/com/google/bitcoin/core/TransactionBroadcastTest.java

@@ -18,9 +18,9 @@
 package com.google.bitcoin.core;
 
 import com.google.bitcoin.params.UnitTestParams;
-import com.google.bitcoin.store.MemoryBlockStore;
 import com.google.bitcoin.testing.FakeTxBuilder;
 import com.google.bitcoin.testing.InboundMessageQueuer;
+import com.google.bitcoin.testing.TestWithPeerGroup;
 import com.google.bitcoin.utils.Threading;
 import com.google.common.util.concurrent.ListenableFuture;
 import org.junit.After;
@@ -35,7 +35,6 @@ import java.util.Random;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 import static org.junit.Assert.*;
-import static org.junit.Assert.assertEquals;
 
 @RunWith(value = Parameterized.class)
 public class TransactionBroadcastTest extends TestWithPeerGroup {
@@ -55,8 +54,7 @@ public class TransactionBroadcastTest extends TestWithPeerGroup {
     @Before
     public void setUp() throws Exception {
         Utils.setMockClock(); // Use mock clock
-        super.setUp(new MemoryBlockStore(UnitTestParams.get()));
+        super.setUp();
-        peerGroup.addWallet(wallet);
         // Fix the random permutation that TransactionBroadcast uses to shuffle the peers.
         TransactionBroadcast.random = new Random(0);
         peerGroup.setMinBroadcastConnections(2);
@@ -64,11 +62,10 @@ public class TransactionBroadcastTest extends TestWithPeerGroup {
         peerGroup.awaitRunning();
     }
 
+    @Override
     @After
-    public void tearDown() throws Exception {
+    public void tearDown() {
         super.tearDown();
-        peerGroup.stopAsync();
-        peerGroup.awaitTerminated();
     }
 
     @Test