Add main-specific node state

2 files changed, 21 insertions, 7 deletions

diff --git a/src/test/DoS_tests.cpp b/src/test/DoS_tests.cpp
index f0fb84bc54..fbca09b4dc 100644
--- a/src/test/DoS_tests.cpp
+++ b/src/test/DoS_tests.cpp
@@ -21,6 +21,7 @@
 // Tests this internal-to-main.cpp method:
 extern bool AddOrphanTx(const CTransaction& tx);
 extern unsigned int LimitOrphanTxSize(unsigned int nMaxOrphans);
+extern void Misbehaving(NodeId nodeid, int howmuch);
 extern std::map<uint256, CTransaction> mapOrphanTransactions;
 extern std::map<uint256, std::set<uint256> > mapOrphanTransactionsByPrev;
 
@@ -38,16 +39,21 @@ BOOST_AUTO_TEST_CASE(DoS_banning)
     CNode::ClearBanned();
     CAddress addr1(ip(0xa0b0c001));
     CNode dummyNode1(INVALID_SOCKET, addr1, "", true);
-    dummyNode1.Misbehaving(100); // Should get banned
+    dummyNode1.nVersion = 1;
+    Misbehaving(dummyNode1.GetId(), 100); // Should get banned
+    SendMessages(&dummyNode1, false);
     BOOST_CHECK(CNode::IsBanned(addr1));
     BOOST_CHECK(!CNode::IsBanned(ip(0xa0b0c001|0x0000ff00))); // Different IP, not banned
 
     CAddress addr2(ip(0xa0b0c002));
     CNode dummyNode2(INVALID_SOCKET, addr2, "", true);
-    dummyNode2.Misbehaving(50);
+    dummyNode2.nVersion = 1;
+    Misbehaving(dummyNode2.GetId(), 50);
+    SendMessages(&dummyNode2, false);
     BOOST_CHECK(!CNode::IsBanned(addr2)); // 2 not banned yet...
     BOOST_CHECK(CNode::IsBanned(addr1));  // ... but 1 still should be
-    dummyNode2.Misbehaving(50);
+    Misbehaving(dummyNode2.GetId(), 50);
+    SendMessages(&dummyNode2, false);
     BOOST_CHECK(CNode::IsBanned(addr2));
 }
 
@@ -57,11 +63,15 @@ BOOST_AUTO_TEST_CASE(DoS_banscore)
     mapArgs["-banscore"] = "111"; // because 11 is my favorite number
     CAddress addr1(ip(0xa0b0c001));
     CNode dummyNode1(INVALID_SOCKET, addr1, "", true);
-    dummyNode1.Misbehaving(100);
+    dummyNode1.nVersion = 1;
+    Misbehaving(dummyNode1.GetId(), 100);
+    SendMessages(&dummyNode1, false);
     BOOST_CHECK(!CNode::IsBanned(addr1));
-    dummyNode1.Misbehaving(10);
+    Misbehaving(dummyNode1.GetId(), 10);
+    SendMessages(&dummyNode1, false);
     BOOST_CHECK(!CNode::IsBanned(addr1));
-    dummyNode1.Misbehaving(1);
+    Misbehaving(dummyNode1.GetId(), 1);
+    SendMessages(&dummyNode1, false);
     BOOST_CHECK(CNode::IsBanned(addr1));
     mapArgs.erase("-banscore");
 }
@@ -74,8 +84,10 @@ BOOST_AUTO_TEST_CASE(DoS_bantime)
 
     CAddress addr(ip(0xa0b0c001));
     CNode dummyNode(INVALID_SOCKET, addr, "", true);
+    dummyNode.nVersion = 1;
 
-    dummyNode.Misbehaving(100);
+    Misbehaving(dummyNode.GetId(), 100);
+    SendMessages(&dummyNode, false);
     BOOST_CHECK(CNode::IsBanned(addr));
 
     SetMockTime(nStartTime+60*60);
diff --git a/src/test/test_bitcoin.cpp b/src/test/test_bitcoin.cpp
index a804ff3803..a4592fe803 100644
--- a/src/test/test_bitcoin.cpp
+++ b/src/test/test_bitcoin.cpp
@@ -47,11 +47,13 @@ struct TestingSetup {
         nScriptCheckThreads = 3;
         for (int i=0; i < nScriptCheckThreads-1; i++)
             threadGroup.create_thread(&ThreadScriptCheck);
+        RegisterNodeSignals(GetNodeSignals());
     }
     ~TestingSetup()
     {
         threadGroup.interrupt_all();
         threadGroup.join_all();
+        UnregisterNodeSignals(GetNodeSignals());
 #ifdef ENABLE_WALLET
         delete pwalletMain;
         pwalletMain = NULL;