Address review comments

beacon_node/network/src/sync/manager.rs

@@ -1225,10 +1225,7 @@ impl<T: BeaconChainTypes> SyncManager<T> {
             // custody_by_range accumulates the results of multiple data_columns_by_range requests
             // returning a bigger list of data columns across all the column indices this node has
             // to custody
-            if let Some(result) =
+            if let Some(result) = self.network.on_custody_by_range_response(id, peer_id, resp) {
-                self.network
-                    .on_custody_by_range_response(id.parent_request_id, id, peer_id, resp)
-            {
                 self.on_custody_by_range_result(id.parent_request_id, result);
             }
         }

beacon_node/network/src/sync/network_context.rs

@@ -1376,14 +1376,13 @@ impl<T: BeaconChainTypes> SyncNetworkContext<T> {
     #[allow(clippy::type_complexity)]
     pub fn on_custody_by_range_response(
         &mut self,
-        id: CustodyByRangeRequestId,
         req_id: DataColumnsByRangeRequestId,
         peer_id: PeerId,
         resp: RpcResponseResult<DataColumnSidecarList<T::EthSpec>>,
     ) -> Option<CustodyRequestResult<T::EthSpec>> {
         // Note: need to remove the request to borrow self again below. Otherwise we can't
         // do nested requests
-        let Some(mut request) = self.custody_by_range_requests.remove(&id) else {
+        let Some(mut request) = self.custody_by_range_requests.remove(&id.parent_request_id) else {
             metrics::inc_counter_vec(
                 &metrics::SYNC_UNKNOWN_NETWORK_REQUESTS,
                 &["custody_by_range"],
@@ -1396,7 +1395,7 @@ impl<T: BeaconChainTypes> SyncNetworkContext<T> {
             .map_err(Into::<RpcResponseError>::into)
             .transpose();
 
-        self.handle_custody_by_range_result(id, request, result)
+        self.handle_custody_by_range_result(id.parent_request_id, request, result)
     }
 
     fn handle_custody_by_range_result(

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

@@ -368,16 +368,19 @@ impl TestRig {
         self.expect_empty_network();
     }
 
-    // Don't make pub, use `add_connected_peer_testing_only`
+    // Note: prefer to use `add_connected_peer_testing_only`. This is currently extensively used in
+    // lookup tests. We should consolidate this "add peer" methods in a future refactor
     fn new_connected_peer(&mut self) -> PeerId {
         self.add_connected_peer_testing_only(false)
     }
 
-    // Don't make pub, use `add_connected_peer_testing_only`
+    // Note: prefer to use `add_connected_peer_testing_only`. This is currently extensively used in
+    // lookup tests. We should consolidate this "add peer" methods in a future refactor
     fn new_connected_supernode_peer(&mut self) -> PeerId {
         self.add_connected_peer_testing_only(true)
     }
 
+    /// Add a random connected peer that is not known by the sync module
     pub fn add_connected_peer_testing_only(&mut self, supernode: bool) -> PeerId {
         let key = self.determinstic_key();
         let peer_id = self
@@ -401,6 +404,7 @@ impl TestRig {
         peer_id
     }
 
+    /// Add a random connected peer + add it to sync with a specific remote Status
     pub fn add_sync_peer(&mut self, supernode: bool, remote_info: SyncInfo) -> PeerId {
         let peer_id = self.add_connected_peer_testing_only(supernode);
         self.send_sync_message(SyncMessage::AddPeer(peer_id, remote_info));
@@ -887,7 +891,7 @@ impl TestRig {
         }
     }
 
-    // Find, not pop
+    /// Similar to `pop_received_network_events` but finds matching events without removing them.
     pub fn filter_received_network_events<T, F: Fn(&NetworkMessage<E>) -> Option<T>>(
         &mut self,
         predicate_transform: F,
@@ -1149,15 +1153,10 @@ impl TestRig {
     }
 
     pub fn expect_no_penalty_for_anyone(&mut self) {
-        self.drain_network_rx();
+        let downscore_events = self.filter_received_network_events(|ev| match ev {
-        let downscore_events = self
+            NetworkMessage::ReportPeer { peer_id, msg, .. } => Some((peer_id, msg)),
-            .network_rx_queue
+            _ => None,
-            .iter()
+        });
-            .filter_map(|ev| match ev {
-                NetworkMessage::ReportPeer { peer_id, msg, .. } => Some((peer_id, msg)),
-                _ => None,
-            })
-            .collect::<Vec<_>>();
         if !downscore_events.is_empty() {
             panic!("Expected no downscoring events but found: {downscore_events:?}");
         }

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

@@ -51,6 +51,7 @@ enum ByRangeDataRequestIds {
 }
 
 impl ByRangeDataRequestIds {
+    /// If there's a single active request, returns its peer, else panics
     fn peer(&self) -> PeerId {
         match self {
             Self::PreDeneb => panic!("no requests PreDeneb"),