From 58d273db2243b9bed1a16f9cd7f22a275f3273d5 Mon Sep 17 00:00:00 2001
From: dapplion <35266934+dapplion@users.noreply.github.com>
Date: Thu, 9 May 2024 12:22:17 +0900
Subject: [PATCH] Fix tests

---
 .../network/src/sync/block_lookups/single_block_lookup.rs   | 6 +++++-
 beacon_node/network/src/sync/block_lookups/tests.rs         | 6 +++---
 beacon_node/network/src/sync/network_context/custody.rs     | 6 ++++--
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/beacon_node/network/src/sync/block_lookups/single_block_lookup.rs b/beacon_node/network/src/sync/block_lookups/single_block_lookup.rs
index 37bb69374d9..9020aea52e8 100644
--- a/beacon_node/network/src/sync/block_lookups/single_block_lookup.rs
+++ b/beacon_node/network/src/sync/block_lookups/single_block_lookup.rs
@@ -1,3 +1,4 @@
+use super::common::ResponseType;
 use super::{BlockComponent, PeerId, SINGLE_BLOCK_LOOKUP_MAX_ATTEMPTS};
 use crate::sync::block_lookups::common::RequestState;
 use crate::sync::block_lookups::Id;
@@ -150,6 +151,7 @@ impl<T: BeaconChainTypes> SingleBlockLookup<T> {
     ) -> Result<(), LookupRequestError> {
         let id = self.id;
         let awaiting_parent = self.awaiting_parent.is_some();
+        let block_is_processed = self.block_request_state.state.is_processed();
         let request = R::request_state_mut(self);
 
         // Attempt to progress awaiting downloads
@@ -181,7 +183,9 @@ impl<T: BeaconChainTypes> SingleBlockLookup<T> {
         // Otherwise, attempt to progress awaiting processing
         // If this request is awaiting a parent lookup to be processed, do not send for processing.
         // The request will be rejected with unknown parent error.
-        } else if !awaiting_parent {
+        } else if !awaiting_parent
+            && (block_is_processed || matches!(R::response_type(), ResponseType::Block))
+        {
             // maybe_start_processing returns Some if state == AwaitingProcess. This pattern is
             // useful to conditionally access the result data.
             if let Some(result) = request.get_state_mut().maybe_start_processing() {
diff --git a/beacon_node/network/src/sync/block_lookups/tests.rs b/beacon_node/network/src/sync/block_lookups/tests.rs
index a678e841d05..f2fa7bd5be6 100644
--- a/beacon_node/network/src/sync/block_lookups/tests.rs
+++ b/beacon_node/network/src/sync/block_lookups/tests.rs
@@ -550,7 +550,7 @@ impl TestRig {
         let peer_id = PeerId::random();
         let slot = block.slot();
         let block_root = block.canonical_root();
-        self.single_lookup_block_response(id, peer_id, Some(block.into()));
+        self.single_lookup_block_response(id, peer_id, Some(block));
         self.single_lookup_block_response(id, peer_id, None);
         // Expect processing and resolve with import
         self.expect_block_process(ResponseType::Block);
@@ -1633,10 +1633,10 @@ fn custody_lookup_happy_path() {
     r.trigger_unknown_block_from_attestation(block_root, peer_id);
     // Should not request blobs
     let id = r.expect_block_lookup_request(block.canonical_root());
+    r.complete_valid_block_request(id, block.into(), true);
     // TODO(das): do not hardcode 4
     let custody_ids = r.expect_only_data_columns_by_root_requests(block_root, 4);
-    r.complete_valid_custody_request(custody_ids, data_columns, true);
-    r.complete_valid_block_request(id, block.into(), false);
+    r.complete_valid_custody_request(custody_ids, data_columns, false);
     r.expect_no_active_lookups();
 }
 
diff --git a/beacon_node/network/src/sync/network_context/custody.rs b/beacon_node/network/src/sync/network_context/custody.rs
index 36c842ae227..72964d3f362 100644
--- a/beacon_node/network/src/sync/network_context/custody.rs
+++ b/beacon_node/network/src/sync/network_context/custody.rs
@@ -44,6 +44,8 @@ pub enum Error {
     NoPeers(ColumnIndex),
 }
 
+type CustodyRequestResult<E> = Result<Option<(Vec<CustodyDataColumn<E>>, PeerGroup)>, Error>;
+
 impl<T: BeaconChainTypes> ActiveCustodyRequest<T> {
     pub(crate) fn new(
         block_root: Hash256,
@@ -80,7 +82,7 @@ impl<T: BeaconChainTypes> ActiveCustodyRequest<T> {
         column_index: ColumnIndex,
         resp: RpcProcessingResult<DataColumnSidecarList<T::EthSpec>>,
         cx: &mut SyncNetworkContext<T>,
-    ) -> Result<Option<(Vec<CustodyDataColumn<T::EthSpec>>, PeerGroup)>, Error> {
+    ) -> CustodyRequestResult<T::EthSpec> {
         // TODO(das): Should downscore peers for verify errors here
 
         let Some(request) = self.column_requests.get_mut(&column_index) else {
@@ -127,7 +129,7 @@ impl<T: BeaconChainTypes> ActiveCustodyRequest<T> {
     pub(crate) fn continue_requests(
         &mut self,
         cx: &mut SyncNetworkContext<T>,
-    ) -> Result<Option<(Vec<CustodyDataColumn<T::EthSpec>>, PeerGroup)>, Error> {
+    ) -> CustodyRequestResult<T::EthSpec> {
         // First check if sampling is completed, by computing `required_successes`
         let mut successes = 0;