From 052845a0da831a6b70327fd64d773fa3b531ba17 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Wed, 21 Aug 2024 14:16:07 +0200
Subject: [PATCH] Fix test

---
 .../src/peer_manager/peerdb.rs                | 47 +++++++++++++++----
 .../network/src/sync/block_lookups/common.rs  | 11 +----
 .../network/src/sync/block_lookups/tests.rs   |  8 ++++
 3 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
index 327ca21a0e4..d602ee2cae3 100644
--- a/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
+++ b/beacon_node/lighthouse_network/src/peer_manager/peerdb.rs
@@ -1,5 +1,6 @@
 use crate::discovery::enr::PEERDAS_CUSTODY_SUBNET_COUNT_ENR_KEY;
 use crate::discovery::CombinedKey;
+use crate::rpc::{MetaData, MetaDataV2};
 use crate::EnrExt;
 use crate::{metrics, multiaddr::Multiaddr, types::Subnet, Enr, Gossipsub, PeerId};
 use peer_info::{ConnectionDirection, PeerConnectionStatus, PeerInfo};
@@ -14,7 +15,7 @@ use std::{
     fmt::Formatter,
 };
 use sync_status::SyncStatus;
-use types::{ChainSpec, EthSpec};
+use types::{ChainSpec, DataColumnSubnetId, EthSpec};
 
 pub mod client;
 pub mod peer_info;
@@ -683,15 +684,20 @@ impl<E: EthSpec> PeerDB<E> {
         let enr_key = CombinedKey::generate_secp256k1();
         let mut enr = Enr::builder().build(&enr_key).unwrap();
         let peer_id = enr.peer_id();
+        let node_id = enr.node_id().raw().into();
 
-        if supernode {
-            enr.insert(
-                PEERDAS_CUSTODY_SUBNET_COUNT_ENR_KEY,
-                &spec.data_column_sidecar_subnet_count,
-                &enr_key,
-            )
-            .expect("u64 can be encoded");
-        }
+        let custody_subnet_count = if supernode {
+            spec.data_column_sidecar_subnet_count
+        } else {
+            spec.custody_requirement
+        };
+
+        enr.insert(
+            PEERDAS_CUSTODY_SUBNET_COUNT_ENR_KEY,
+            &custody_subnet_count,
+            &enr_key,
+        )
+        .expect("u64 can be encoded");
 
         self.update_connection_state(
             &peer_id,
@@ -701,6 +707,29 @@ impl<E: EthSpec> PeerDB<E> {
                 direction: ConnectionDirection::Outgoing,
             },
         );
+        let peer = self.peers.get_mut(&peer_id).expect("peer exists");
+
+        // Need to insert an empty metadata to pass the condition `on_subnet_metadata`
+        peer.set_meta_data(MetaData::V2(MetaDataV2 {
+            seq_number: 0,
+            attnets: <_>::default(),
+            syncnets: <_>::default(),
+        }));
+
+        for subnet in
+            DataColumnSubnetId::compute_custody_columns::<E>(node_id, custody_subnet_count, spec)
+        {
+            // Need to pass the tests on PeerInfo
+            // - on_subnet_metadata: no action, `Subnet::DataColumn` returns true
+            // - on_subnet_gossipsub: subnets field contains subnet
+            // - is_good_gossipsub_peer: score >= 0, which equals default score
+            // peer.set_meta_data(crate::rpc::MetaData::V2(crate::rpc::MetaDataV2 {
+            //     seq_number: 0,
+            //     attnets: <_>::default(),
+            //     syncnets: <_>::default(),
+            // }));
+            peer.insert_subnet(Subnet::DataColumn(subnet.into()));
+        }
 
         peer_id
     }
diff --git a/beacon_node/network/src/sync/block_lookups/common.rs b/beacon_node/network/src/sync/block_lookups/common.rs
index 5fe7d1dca4a..a7be72556e2 100644
--- a/beacon_node/network/src/sync/block_lookups/common.rs
+++ b/beacon_node/network/src/sync/block_lookups/common.rs
@@ -171,15 +171,8 @@ impl<T: BeaconChainTypes> RequestState<T> for CustodyRequestState<T::EthSpec> {
         downloaded_block: Option<Arc<SignedBeaconBlock<T::EthSpec>>>,
         cx: &mut SyncNetworkContext<T>,
     ) -> Result<LookupRequestResult, LookupRequestError> {
-        let x = cx
-            .custody_lookup_request(id, self.block_root, downloaded_block)
-            .map_err(LookupRequestError::SendFailedNetwork);
-        match x.as_ref().unwrap() {
-            LookupRequestResult::RequestSent { .. } => println!("RequestSent"),
-            LookupRequestResult::NoRequestNeeded(s) => println!("NoRequestNeeded {}", s),
-            LookupRequestResult::Pending(s) => println!("Pending {}", s),
-        }
-        x
+        cx.custody_lookup_request(id, self.block_root, downloaded_block)
+            .map_err(LookupRequestError::SendFailedNetwork)
     }
 
     fn send_for_processing(
diff --git a/beacon_node/network/src/sync/block_lookups/tests.rs b/beacon_node/network/src/sync/block_lookups/tests.rs
index a96598eaac7..43b73b6e6df 100644
--- a/beacon_node/network/src/sync/block_lookups/tests.rs
+++ b/beacon_node/network/src/sync/block_lookups/tests.rs
@@ -377,6 +377,14 @@ impl TestRig {
         }
         // One supernode peer to ensure all columns have at least one peer
         self.new_connected_supernode_peer();
+
+        // ensure we have peers in each subnet
+        for i in 0..self.harness.spec.data_column_sidecar_subnet_count {
+            let peers = self
+                .network_globals
+                .custody_peers_for_column(i, &self.harness.spec);
+            println!("peers column {} {:?}", i, peers);
+        }
     }
 
     fn parent_chain_processed_success(