Remove BatchStateSummary

beacon_node/network/src/sync/range_sync/chain.rs

@@ -129,37 +129,13 @@ pub enum ChainSyncingState {
     Syncing,
 }
 
-#[cfg(test)]
-#[derive(Debug, Eq, PartialEq)]
-pub enum BatchStateSummary {
-    Downloading,
-    Processing,
-    AwaitingProcessing,
-    AwaitingValidation,
-    Unexpected(&'static str),
-}
-
 impl<T: BeaconChainTypes> SyncingChain<T> {
-    /// Returns a summary of batch states for assertions in tests.
+    /// Leaks the state of all active batches for assertions in tests.
     #[cfg(test)]
-    pub fn batches_state(&self) -> Vec<(BatchId, BatchStateSummary)> {
+    pub fn batches_state(&self) -> Vec<(BatchId, &BatchState<T::EthSpec>)> {
         self.batches
             .iter()
-            .map(|(id, batch)| {
-                let state = match batch.state() {
-                    // A batch is never left in this state, it's only the initial value
-                    BatchState::AwaitingDownload => {
-                        BatchStateSummary::Unexpected("AwaitingDownload")
-                    }
-                    BatchState::Downloading { .. } => BatchStateSummary::Downloading,
-                    BatchState::AwaitingProcessing { .. } => BatchStateSummary::AwaitingProcessing,
-                    BatchState::Poisoned => BatchStateSummary::Unexpected("Poisoned"),
-                    BatchState::Processing { .. } => BatchStateSummary::Processing,
-                    BatchState::Failed => BatchStateSummary::Unexpected("Failed"),
-                    BatchState::AwaitingValidation { .. } => BatchStateSummary::AwaitingValidation,
-                };
-                (*id, state)
-            })
+            .map(|(id, batch)| (*id, batch.state()))
             .collect()
     }
 

beacon_node/network/src/sync/range_sync/mod.rs

@@ -10,8 +10,6 @@ mod sync_type;
 pub use batch::{
     BatchConfig, BatchInfo, BatchOperationOutcome, BatchPeers, BatchProcessingResult, BatchState,
 };
-#[cfg(test)]
-pub use chain::BatchStateSummary;
 pub use chain::{BatchId, ChainId, EPOCHS_PER_BATCH};
 pub use range::RangeSync;
 pub use sync_type::RangeSyncType;

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

@@ -39,8 +39,6 @@
 //!  Each chain is downloaded in batches of blocks. The batched blocks are processed sequentially
 //!  and further batches are requested as current blocks are being processed.
 
-#[cfg(test)]
-use super::chain::BatchStateSummary;
 use super::chain::{BatchId, ChainId, RemoveChain, SyncingChain};
 use super::chain_collection::{ChainCollection, SyncChainStatus};
 use super::sync_type::RangeSyncType;
@@ -48,6 +46,8 @@ use super::BatchPeers;
 use crate::metrics;
 use crate::status::ToStatusMessage;
 use crate::sync::network_context::{RpcResponseError, SyncNetworkContext};
+#[cfg(test)]
+use crate::sync::range_sync::BatchState;
 use crate::sync::BatchProcessResult;
 use beacon_chain::block_verification_types::RpcBlock;
 use beacon_chain::{BeaconChain, BeaconChainTypes};
@@ -107,7 +107,7 @@ where
     }
 
     #[cfg(test)]
-    pub(crate) fn batches_state(&self) -> Vec<(ChainId, BatchId, BatchStateSummary)> {
+    pub(crate) fn batches_state(&self) -> Vec<(ChainId, BatchId, &BatchState<T::EthSpec>)> {
         self.chains
             .iter()
             .flat_map(|chain| {

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

@@ -3,7 +3,7 @@ use crate::network_beacon_processor::ChainSegmentProcessId;
 use crate::status::ToStatusMessage;
 use crate::sync::manager::SLOT_IMPORT_TOLERANCE;
 use crate::sync::network_context::{BlockComponentsByRangeRequestStep, RangeRequestId};
-use crate::sync::range_sync::{BatchId, BatchStateSummary, RangeSyncType};
+use crate::sync::range_sync::{BatchId, BatchState, RangeSyncType};
 use crate::sync::{ChainId, SyncMessage};
 use beacon_chain::data_column_verification::CustodyDataColumn;
 use beacon_chain::test_utils::{test_spec, AttestationStrategy, BlockStrategy};
@@ -298,7 +298,7 @@ impl TestRig {
         self.sync_manager.network().network_globals().sync_state()
     }
 
-    fn get_batch_states(&mut self) -> Vec<(ChainId, BatchId, BatchStateSummary)> {
+    fn get_batch_states(&mut self) -> Vec<(ChainId, BatchId, &BatchState<E>)> {
         self.sync_manager.range_sync().batches_state()
     }
 
@@ -382,27 +382,39 @@ impl TestRig {
         }
     }
 
-    fn expect_all_batches_in_state(&mut self, states: &[BatchStateSummary]) {
+    fn expect_all_batches_in_state<F: Fn(&BatchState<E>) -> bool>(
+        &mut self,
+        predicate: F,
+        expected_state: &'static str,
+    ) {
         let batches = self.get_batch_states();
         if batches.is_empty() {
             panic!("no batches");
         }
-        for batch in &batches {
-            if !states.contains(&batch.2) {
-                panic!("batch {batch:?} not in state {states:?}. Batches: {batches:?}");
+        for (chain_id, batch_id, state) in &batches {
+            if !predicate(state) {
+                panic!("batch {chain_id} {batch_id} not in state {expected_state}, {state}");
             }
         }
     }
 
     fn expect_all_batches_downloading(&mut self) {
-        self.expect_all_batches_in_state(&[BatchStateSummary::Downloading]);
+        self.expect_all_batches_in_state(
+            |state| matches!(state, BatchState::Downloading { .. }),
+            "Downloading",
+        );
     }
 
     fn expect_all_batches_processing_or_awaiting(&mut self) {
-        self.expect_all_batches_in_state(&[
-            BatchStateSummary::Processing,
-            BatchStateSummary::AwaitingProcessing,
-        ]);
+        self.expect_all_batches_in_state(
+            |state| {
+                matches!(
+                    state,
+                    BatchState::Processing { .. } | BatchState::AwaitingProcessing { .. }
+                )
+            },
+            "Processing or AwaitingProcessing",
+        );
     }
 
     fn update_execution_engine_state(&mut self, state: EngineState) {