From 72e30e8e03f880eba4bd1c3fc18b5558d8cef680 Mon Sep 17 00:00:00 2001
From: Jon Atack <jon@atack.com>
Date: Sun, 21 Feb 2021 16:36:49 +0100
Subject: [PATCH] Add unit tests for ProtectEvictionCandidatesByRatio()

Thank you to Vasil Dimov (vasild) for the suggestion to use std::unordered_set
rather than std::vector for the IsProtected() peer id arguments.
---
 src/test/net_peer_eviction_tests.cpp | 104 +++++++++++++++++++++++++--
 1 file changed, 97 insertions(+), 7 deletions(-)

diff --git a/src/test/net_peer_eviction_tests.cpp b/src/test/net_peer_eviction_tests.cpp
index 290a6b7eaea..418f5a4f713 100644
--- a/src/test/net_peer_eviction_tests.cpp
+++ b/src/test/net_peer_eviction_tests.cpp
@@ -15,6 +15,11 @@
 
 BOOST_FIXTURE_TEST_SUITE(net_peer_eviction_tests, BasicTestingSetup)
 
+namespace {
+constexpr int NODE_EVICTION_TEST_ROUNDS{10};
+constexpr int NODE_EVICTION_TEST_UP_TO_N_NODES{200};
+} // namespace
+
 std::vector<NodeEvictionCandidate> GetRandomNodeEvictionCandidates(const int n_candidates, FastRandomContext& random_context)
 {
     std::vector<NodeEvictionCandidate> candidates;
@@ -36,6 +41,98 @@ std::vector<NodeEvictionCandidate> GetRandomNodeEvictionCandidates(const int n_c
     return candidates;
 }
 
+// Create `num_peers` random nodes, apply setup function `candidate_setup_fn`,
+// call ProtectEvictionCandidatesByRatio() to apply protection logic, and then
+// return true if all of `protected_peer_ids` and none of `unprotected_peer_ids`
+// are protected from eviction, i.e. removed from the eviction candidates.
+bool IsProtected(int num_peers,
+                 std::function<void(NodeEvictionCandidate&)> candidate_setup_fn,
+                 const std::unordered_set<NodeId>& protected_peer_ids,
+                 const std::unordered_set<NodeId>& unprotected_peer_ids,
+                 FastRandomContext& random_context)
+{
+    std::vector<NodeEvictionCandidate> candidates{GetRandomNodeEvictionCandidates(num_peers, random_context)};
+    for (NodeEvictionCandidate& candidate : candidates) {
+        candidate_setup_fn(candidate);
+    }
+    Shuffle(candidates.begin(), candidates.end(), random_context);
+
+    const size_t size{candidates.size()};
+    const size_t expected{size - size / 2}; // Expect half the candidates will be protected.
+    ProtectEvictionCandidatesByRatio(candidates);
+    BOOST_CHECK_EQUAL(candidates.size(), expected);
+
+    size_t unprotected_count{0};
+    for (const NodeEvictionCandidate& candidate : candidates) {
+        if (protected_peer_ids.count(candidate.id)) {
+            // this peer should have been removed from the eviction candidates
+            BOOST_TEST_MESSAGE(strprintf("expected candidate to be protected: %d", candidate.id));
+            return false;
+        }
+        if (unprotected_peer_ids.count(candidate.id)) {
+            // this peer remains in the eviction candidates, as expected
+            ++unprotected_count;
+        }
+    }
+
+    const bool is_protected{unprotected_count == unprotected_peer_ids.size()};
+    if (!is_protected) {
+        BOOST_TEST_MESSAGE(strprintf("unprotected: expected %d, actual %d",
+                                     unprotected_peer_ids.size(), unprotected_count));
+    }
+    return is_protected;
+}
+
+BOOST_AUTO_TEST_CASE(peer_protection_test)
+{
+    FastRandomContext random_context{true};
+    int num_peers{12};
+
+    // Expect half of the peers with greatest uptime (the lowest nTimeConnected)
+    // to be protected from eviction.
+    BOOST_CHECK(IsProtected(
+        num_peers, [](NodeEvictionCandidate& c) {
+            c.nTimeConnected = c.id;
+            c.m_is_local = false;
+        },
+        /* protected_peer_ids */ {0, 1, 2, 3, 4, 5},
+        /* unprotected_peer_ids */ {6, 7, 8, 9, 10, 11},
+        random_context));
+
+    // Verify in the opposite direction.
+    BOOST_CHECK(IsProtected(
+        num_peers, [num_peers](NodeEvictionCandidate& c) {
+            c.nTimeConnected = num_peers - c.id;
+            c.m_is_local = false;
+        },
+        /* protected_peer_ids */ {6, 7, 8, 9, 10, 11},
+        /* unprotected_peer_ids */ {0, 1, 2, 3, 4, 5},
+        random_context));
+
+    // Test protection of localhost peers...
+
+    // Expect 1/4 localhost peers to be protected from eviction,
+    // independently of other characteristics.
+    BOOST_CHECK(IsProtected(
+        num_peers, [](NodeEvictionCandidate& c) {
+            c.m_is_local = (c.id == 1 || c.id == 9 || c.id == 11);
+        },
+        /* protected_peer_ids */ {1, 9, 11},
+        /* unprotected_peer_ids */ {},
+        random_context));
+
+    // Expect 1/4 localhost peers and 1/4 of the others to be protected
+    // from eviction, sorted by longest uptime (lowest nTimeConnected).
+    BOOST_CHECK(IsProtected(
+        num_peers, [](NodeEvictionCandidate& c) {
+            c.nTimeConnected = c.id;
+            c.m_is_local = (c.id > 6);
+        },
+        /* protected_peer_ids */ {0, 1, 2, 7, 8, 9},
+        /* unprotected_peer_ids */ {3, 4, 5, 6, 10, 11},
+        random_context));
+}
+
 // Returns true if any of the node ids in node_ids are selected for eviction.
 bool IsEvicted(std::vector<NodeEvictionCandidate> candidates, const std::unordered_set<NodeId>& node_ids, FastRandomContext& random_context)
 {
@@ -59,11 +156,6 @@ bool IsEvicted(const int number_of_nodes, std::function<void(NodeEvictionCandida
     return IsEvicted(candidates, node_ids, random_context);
 }
 
-namespace {
-constexpr int NODE_EVICTION_TEST_ROUNDS{10};
-constexpr int NODE_EVICTION_TEST_UP_TO_N_NODES{200};
-} // namespace
-
 BOOST_AUTO_TEST_CASE(peer_eviction_test)
 {
     FastRandomContext random_context{true};
@@ -150,8 +242,6 @@ BOOST_AUTO_TEST_CASE(peer_eviction_test)
             }
 
             // Cases left to test:
-            // * "Protect the half of the remaining nodes which have been connected the longest. [...]"
-            // * "Pick out up to 1/4 peers that are localhost, sorted by longest uptime. [...]"
             // * "If any remaining peers are preferred for eviction consider only them. [...]"
             // * "Identify the network group with the most connections and youngest member. [...]"
         }