From 9426f6e9626bff195ef7a3b9bb5d2288bddc7d55 Mon Sep 17 00:00:00 2001
From: Federico Gimenez <federico.gimenez@gmail.com>
Date: Wed, 7 Aug 2024 11:19:15 +0200
Subject: [PATCH] test: live sync, FCU extends canon chain

---
 crates/engine/tree/src/tree/mod.rs | 72 ++++++++++++++++++++++--------
 1 file changed, 53 insertions(+), 19 deletions(-)

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index 5796756596e6d..556f0f6d09c56 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -2125,8 +2125,7 @@ mod tests {
             ));
 
             let response = rx.await.unwrap().unwrap().await.unwrap();
-            let fcu_status = fcu_status.into();
-            match fcu_status {
+            match fcu_status.into() {
                 ForkchoiceStatus::Valid => assert!(response.payload_status.is_valid()),
                 ForkchoiceStatus::Syncing => assert!(response.payload_status.is_syncing()),
                 ForkchoiceStatus::Invalid => assert!(response.payload_status.is_invalid()),
@@ -2232,6 +2231,10 @@ mod tests {
             }
             self.extend_execution_outcome(execution_outcomes);
         }
+
+        fn check_canon_head(&self, head_hash: B256) {
+            assert_eq!(self.tree.state.tree_state.canonical_head().hash, head_hash);
+        }
     }
 
     #[tokio::test]
@@ -2640,27 +2643,25 @@ mod tests {
         test_harness = test_harness.with_blocks(main_chain.clone());
 
         let fork_chain = test_harness.block_builder.create_fork(main_chain[2].block(), 3);
+        let fork_chain_last_hash = fork_chain.last().unwrap().hash();
 
         // add fork blocks to the tree
         for block in &fork_chain {
             test_harness.insert_block(block.clone()).unwrap();
         }
 
-        test_harness.send_fcu(fork_chain.last().unwrap().hash(), ForkchoiceStatus::Valid).await;
+        test_harness.send_fcu(fork_chain_last_hash, ForkchoiceStatus::Valid).await;
 
         // check for ForkBlockAdded events, we expect fork_chain.len() blocks added
         test_harness.check_fork_chain_insertion(fork_chain.clone()).await;
 
         // check for CanonicalChainCommitted event
-        test_harness.check_canon_commit(fork_chain.last().unwrap().hash()).await;
+        test_harness.check_canon_commit(fork_chain_last_hash).await;
 
-        test_harness.check_fcu(fork_chain.last().unwrap().hash(), ForkchoiceStatus::Valid).await;
+        test_harness.check_fcu(fork_chain_last_hash, ForkchoiceStatus::Valid).await;
 
         // new head is the tip of the fork chain
-        assert_eq!(
-            test_harness.tree.state.tree_state.canonical_head().hash,
-            fork_chain.last().unwrap().hash()
-        );
+        test_harness.check_canon_head(fork_chain_last_hash);
     }
 
     #[tokio::test]
@@ -2836,17 +2837,50 @@ mod tests {
         test_harness.check_fork_chain_insertion(remaining).await;
 
         // check canonical chain committed event with the hash of the latest block
-        let event = test_harness.from_tree_rx.recv().await.unwrap();
-        match event {
-            EngineApiEvent::BeaconConsensus(
-                BeaconConsensusEngineEvent::CanonicalChainCommitted(header, ..),
-            ) => {
-                assert_eq!(header.hash(), main_chain_last_hash);
-            }
-            _ => panic!("Unexpected event: {:#?}", event),
-        }
+        test_harness.check_canon_commit(main_chain_last_hash).await;
 
         // new head is the tip of the main chain
-        assert_eq!(test_harness.tree.state.tree_state.canonical_head().hash, main_chain_last_hash);
+        test_harness.check_canon_head(main_chain_last_hash);
+    }
+
+    #[tokio::test]
+    async fn test_engine_tree_live_sync_fcu_extends_canon_chain() {
+        reth_tracing::init_test_tracing();
+
+        let chain_spec = MAINNET.clone();
+        let mut test_harness = TestHarness::new(chain_spec.clone());
+
+        // create base chain and setup test harness with it
+        let base_chain: Vec<_> = test_harness.block_builder.get_executed_blocks(0..1).collect();
+        test_harness = test_harness.with_blocks(base_chain.clone());
+
+        // fcu to the tip of base chain
+        test_harness
+            .fcu_to(base_chain.last().unwrap().block().hash(), ForkchoiceStatus::Valid)
+            .await;
+
+        // create main chain, extension of base chain
+        let main_chain = test_harness.block_builder.create_fork(base_chain[0].block(), 10);
+        // determine target in the middle of main hain
+        let target = main_chain.get(5).unwrap();
+        let target_hash = target.hash();
+        let main_last = main_chain.last().unwrap();
+        let main_last_hash = main_last.hash();
+
+        // insert main chain
+        test_harness.insert_chain(main_chain).await;
+
+        // send fcu to target
+        test_harness.send_fcu(target_hash, ForkchoiceStatus::Valid).await;
+
+        test_harness.check_canon_commit(target_hash).await;
+        test_harness.check_fcu(target_hash, ForkchoiceStatus::Valid).await;
+
+        // send fcu to main tip
+        test_harness.send_fcu(main_last_hash, ForkchoiceStatus::Valid).await;
+
+        test_harness.check_canon_commit(main_last_hash).await;
+        test_harness.check_fcu(main_last_hash, ForkchoiceStatus::Valid).await;
+        test_harness.check_canon_head(main_last_hash);
     }
 }