From 637bb6da368b87711005b909f451f94909400092 Mon Sep 17 00:00:00 2001
From: Vasil Dimov <vd@FreeBSD.org>
Date: Wed, 20 Jan 2021 11:40:01 +0100
Subject: [PATCH] test: also check disconnect in
 denialofservice_tests/peer_discouragement

Use `CConnmanTest` instead of `CConnman` and add the nodes to it
so that their `fDisconnect` flag is set during disconnection.
---
 src/test/denialofservice_tests.cpp | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/src/test/denialofservice_tests.cpp b/src/test/denialofservice_tests.cpp
index 1769664a6ca..f556585b42f 100644
--- a/src/test/denialofservice_tests.cpp
+++ b/src/test/denialofservice_tests.cpp
@@ -210,7 +210,7 @@ BOOST_AUTO_TEST_CASE(peer_discouragement)
 {
     const CChainParams& chainparams = Params();
     auto banman = std::make_unique<BanMan>(GetDataDir() / "banlist.dat", nullptr, DEFAULT_MISBEHAVING_BANTIME);
-    auto connman = std::make_unique<CConnman>(0x1337, 0x1337, *m_node.addrman);
+    auto connman = std::make_unique<CConnmanTest>(0x1337, 0x1337, *m_node.addrman);
     auto peerLogic = PeerManager::make(chainparams, *connman, *m_node.addrman, banman.get(),
                                        *m_node.scheduler, *m_node.chainman, *m_node.mempool, false);
 
@@ -228,12 +228,14 @@ BOOST_AUTO_TEST_CASE(peer_discouragement)
     nodes[0]->SetCommonVersion(PROTOCOL_VERSION);
     peerLogic->InitializeNode(nodes[0]);
     nodes[0]->fSuccessfullyConnected = true;
+    connman->AddNode(*nodes[0]);
     peerLogic->Misbehaving(nodes[0]->GetId(), DISCOURAGEMENT_THRESHOLD, /* message */ ""); // Should be discouraged
     {
         LOCK(nodes[0]->cs_sendProcessing);
         BOOST_CHECK(peerLogic->SendMessages(nodes[0]));
     }
     BOOST_CHECK(banman->IsDiscouraged(addr[0]));
+    BOOST_CHECK(nodes[0]->fDisconnect);
     BOOST_CHECK(!banman->IsDiscouraged(other_addr)); // Different address, not discouraged
 
     nodes[1] = new CNode{id++, NODE_NETWORK, INVALID_SOCKET, addr[1], /* nKeyedNetGroupIn */ 1,
@@ -242,26 +244,33 @@ BOOST_AUTO_TEST_CASE(peer_discouragement)
     nodes[1]->SetCommonVersion(PROTOCOL_VERSION);
     peerLogic->InitializeNode(nodes[1]);
     nodes[1]->fSuccessfullyConnected = true;
+    connman->AddNode(*nodes[1]);
     peerLogic->Misbehaving(nodes[1]->GetId(), DISCOURAGEMENT_THRESHOLD - 1, /* message */ "");
     {
         LOCK(nodes[1]->cs_sendProcessing);
         BOOST_CHECK(peerLogic->SendMessages(nodes[1]));
     }
-    BOOST_CHECK(!banman->IsDiscouraged(addr[1])); // [1] not discouraged yet...
-    BOOST_CHECK(banman->IsDiscouraged(addr[0]));  // ... but [0] still should be
+    // [0] is still discouraged/disconnected.
+    BOOST_CHECK(banman->IsDiscouraged(addr[0]));
+    BOOST_CHECK(nodes[0]->fDisconnect);
+    // [1] is not discouraged/disconnected yet.
+    BOOST_CHECK(!banman->IsDiscouraged(addr[1]));
+    BOOST_CHECK(!nodes[1]->fDisconnect);
     peerLogic->Misbehaving(nodes[1]->GetId(), 1, /* message */ ""); // [1] reaches discouragement threshold
     {
         LOCK(nodes[1]->cs_sendProcessing);
         BOOST_CHECK(peerLogic->SendMessages(nodes[1]));
     }
-    // Expect both [0] and [1] to be discouraged now.
+    // Expect both [0] and [1] to be discouraged/disconnected now.
     BOOST_CHECK(banman->IsDiscouraged(addr[0]));
+    BOOST_CHECK(nodes[0]->fDisconnect);
     BOOST_CHECK(banman->IsDiscouraged(addr[1]));
+    BOOST_CHECK(nodes[1]->fDisconnect);
 
     for (CNode* node : nodes) {
         peerLogic->FinalizeNode(*node);
-        delete node;
     }
+    connman->ClearNodes();
 }
 
 BOOST_AUTO_TEST_CASE(DoS_bantime)