Set custody subnets on tests (#6382)

* Set custody subnets on tests

* Enable sampling_with_retries test

* Add another supernode to ensure the node can retry

beacon_node/lighthouse_network/src/peer_manager/peerdb.rs

@@ -1,5 +1,5 @@
 use crate::discovery::enr::PEERDAS_CUSTODY_SUBNET_COUNT_ENR_KEY;
-use crate::discovery::CombinedKey;
+use crate::discovery::{peer_id_to_node_id, CombinedKey};
 use crate::{metrics, multiaddr::Multiaddr, types::Subnet, Enr, EnrExt, Gossipsub, PeerId};
 use peer_info::{ConnectionDirection, PeerConnectionStatus, PeerInfo};
 use rand::seq::SliceRandom;
@@ -723,6 +723,17 @@ impl<E: EthSpec> PeerDB<E> {
                 .map(|csc| csc.into())
                 .collect();
             peer_info.set_custody_subnets(all_subnets);
+        } else {
+            let peer_info = self.peers.get_mut(&peer_id).expect("peer exists");
+            let node_id = peer_id_to_node_id(&peer_id).expect("convert peer_id to node_id");
+            let subnets = DataColumnSubnetId::compute_custody_subnets::<E>(
+                node_id.raw(),
+                spec.custody_requirement,
+                spec,
+            )
+            .expect("should compute custody subnets")
+            .collect();
+            peer_info.set_custody_subnets(subnets);
         }
 
         peer_id

beacon_node/network/src/sync/block_lookups/tests.rs

@@ -1972,12 +1972,13 @@ fn sampling_happy_path() {
 }
 
 #[test]
-#[ignore] // Ignoring due to flakiness https://github.com/sigp/lighthouse/issues/6319
 fn sampling_with_retries() {
     let Some(mut r) = TestRig::test_setup_after_peerdas() else {
         return;
     };
     r.new_connected_peers_for_peerdas();
+    // Add another supernode to ensure that the node can retry.
+    r.new_connected_supernode_peer();
     let (block, data_columns) = r.rand_block_and_data_columns();
     let block_root = block.canonical_root();
     r.trigger_sample_block(block_root, block.slot());