From 5459dea74e4c43bc287f389f420e991b1eb290f9 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Mon, 13 May 2024 23:04:43 +0300
Subject: [PATCH] Make sampling tests deterministic

---
 .../src/peer_manager/peerdb.rs                | 20 ++++++++++++--
 .../network/src/sync/block_lookups/tests.rs   | 26 ++++++++++++++-----
 2 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
index c3e77ae225e..76dbd483b57 100644
--- a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
+++ b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
@@ -1,9 +1,11 @@
+use crate::discovery::enr::PEERDAS_CUSTODY_SUBNET_COUNT_ENR_KEY;
 use crate::discovery::CombinedKey;
 use crate::{metrics, multiaddr::Multiaddr, types::Subnet, Enr, Gossipsub, PeerId};
 use peer_info::{ConnectionDirection, PeerConnectionStatus, PeerInfo};
 use rand::seq::SliceRandom;
 use score::{PeerAction, ReportSource, Score, ScoreState};
 use slog::{crit, debug, error, trace, warn};
+use ssz::Encode;
 use std::net::IpAddr;
 use std::time::Instant;
 use std::{cmp::Ordering, fmt::Display};
@@ -673,9 +675,23 @@ impl<E: EthSpec> PeerDB<E> {
     }
 
     /// Updates the connection state. MUST ONLY BE USED IN TESTS.
-    pub fn __add_connected_peer_testing_only(&mut self, peer_id: &PeerId) -> Option<BanOperation> {
+    pub fn __add_connected_peer_testing_only(
+        &mut self,
+        peer_id: &PeerId,
+        supernode: bool,
+    ) -> Option<BanOperation> {
         let enr_key = CombinedKey::generate_secp256k1();
-        let enr = Enr::builder().build(&enr_key).unwrap();
+        let mut enr = Enr::builder().build(&enr_key).unwrap();
+
+        if supernode {
+            enr.insert(
+                PEERDAS_CUSTODY_SUBNET_COUNT_ENR_KEY,
+                &(E::data_column_subnet_count() as u64).as_ssz_bytes(),
+                &enr_key,
+            )
+            .expect("u64 can be encoded");
+        }
+
         self.update_connection_state(
             peer_id,
             NewConnectionState::Connected {
diff --git a/beacon_node/network/src/sync/block_lookups/tests.rs b/beacon_node/network/src/sync/block_lookups/tests.rs
index 63f00138d67..388c0c09e32 100644
--- a/beacon_node/network/src/sync/block_lookups/tests.rs
+++ b/beacon_node/network/src/sync/block_lookups/tests.rs
@@ -356,12 +356,26 @@ impl TestRig {
         self.network_globals
             .peers
             .write()
-            .__add_connected_peer_testing_only(&peer_id);
+            .__add_connected_peer_testing_only(&peer_id, false);
         peer_id
     }
 
-    fn new_connected_peers(&mut self, count: usize) -> Vec<PeerId> {
-        (0..count).map(|_| self.new_connected_peer()).collect()
+    fn new_connected_supernode_peer(&mut self) -> PeerId {
+        let peer_id = PeerId::random();
+        self.network_globals
+            .peers
+            .write()
+            .__add_connected_peer_testing_only(&peer_id, true);
+        peer_id
+    }
+
+    fn new_connected_peers_for_peerdas(&mut self) {
+        // Enough sampling peers with few columns
+        for _ in 0..100 {
+            self.new_connected_peer();
+        }
+        // One supernode peer to ensure all columns have at least one peer
+        self.new_connected_supernode_peer();
     }
 
     fn parent_chain_processed_success(
@@ -1584,7 +1598,7 @@ fn sampling_happy_path() {
     let Some(mut r) = TestRig::test_setup_after_peerdas() else {
         return;
     };
-    r.new_connected_peers(100); // Add enough sampling peers
+    r.new_connected_peers_for_peerdas();
     let (block, data_columns) = r.rand_block_and_data_columns();
     let block_root = block.canonical_root();
     r.trigger_sample_block(block_root, block.slot());
@@ -1601,7 +1615,7 @@ fn sampling_with_retries() {
     let Some(mut r) = TestRig::test_setup_after_peerdas() else {
         return;
     };
-    r.new_connected_peers(100); // Add enough sampling peers
+    r.new_connected_peers_for_peerdas();
     let (block, data_columns) = r.rand_block_and_data_columns();
     let block_root = block.canonical_root();
     r.trigger_sample_block(block_root, block.slot());
@@ -1621,7 +1635,7 @@ fn custody_lookup_happy_path() {
     let Some(mut r) = TestRig::test_setup_after_peerdas() else {
         return;
     };
-    r.new_connected_peers(100); // Add enough sampling peers
+    r.new_connected_peers_for_peerdas();
     let (block, data_columns) = r.rand_block_and_data_columns();
     let block_root = block.canonical_root();
     let peer_id = r.new_connected_peer();