Fix invalid payload recovery tests

beacon_node/beacon_chain/tests/payload_invalidation.rs

@@ -1821,81 +1821,94 @@ struct InvalidHeadSetup {
 }
 
 impl InvalidHeadSetup {
+    /// This function aims to produce two things:
+    ///
+    /// 1. A chain where the only viable head block has an invalid execution payload.
+    /// 2. A block (`fork_block`) which will become the head of the chain when
+    ///     it is imported.
     async fn new() -> InvalidHeadSetup {
+        let slots_per_epoch = E::slots_per_epoch();
         let mut rig = InvalidPayloadRig::new().enable_attestations();
         rig.move_to_terminal_block();
         rig.import_block(Payload::Valid).await; // Import a valid transition block.
 
-        // Import blocks until the first time the chain finalizes.
+        // Import blocks until the first time the chain finalizes. This avoids
+        // some edge-cases around genesis.
         while rig.cached_head().finalized_checkpoint().epoch == 0 {
             rig.import_block(Payload::Syncing).await;
         }
 
-        let slots_per_epoch = E::slots_per_epoch();
-        let start_slot = rig.cached_head().head_slot() + 1;
-        let mut opt_fork_block = None;
-
-        assert_eq!(start_slot % slots_per_epoch, 1);
-        for i in 0..slots_per_epoch - 1 {
-            let slot = start_slot + i;
-            let slot_offset = slot.as_u64() % slots_per_epoch;
-
-            rig.harness.set_current_slot(slot);
-
-            if slot_offset == slots_per_epoch - 1 {
-                // Optimistic head block right before epoch boundary.
-                let is_valid = Payload::Syncing;
-                rig.import_block_parametric(is_valid, is_valid, Some(slot), |error| {
-                    matches!(
-                        error,
-                        BlockError::ExecutionPayloadError(
-                            ExecutionPayloadError::RejectedByExecutionEngine { .. }
-                        )
-                    )
-                })
-                .await;
-            } else if 3 * slot_offset < 2 * slots_per_epoch {
-                // Valid block in previous epoch.
-                rig.import_block(Payload::Valid).await;
-            } else if slot_offset == slots_per_epoch - 2 {
-                // Fork block one slot prior to invalid head, not applied immediately.
-                let parent_state = rig
-                    .harness
-                    .chain
-                    .state_at_slot(slot - 1, StateSkipConfig::WithStateRoots)
-                    .unwrap();
-                let (fork_block_tuple, _) = rig.harness.make_block(parent_state, slot).await;
-                opt_fork_block = Some(Arc::new(fork_block_tuple.0));
-            } else {
-                // Skipped slot.
-            };
+        // Define a helper function.
+        let chain = rig.harness.chain.clone();
+        let get_unrealized_justified_epoch = move || {
+            chain
+                .canonical_head
+                .fork_choice_read_lock()
+                .unrealized_justified_checkpoint()
+                .epoch
+        };
+
+        // Import more blocks until there is a new and higher unrealized
+        // justified checkpoint.
+        //
+        // The result will be a single chain where the head block has a higher
+        // unrealized justified checkpoint than all other blocks in the chain.
+        let initial_unrealized_justified = get_unrealized_justified_epoch();
+        while get_unrealized_justified_epoch() == initial_unrealized_justified {
+            rig.import_block(Payload::Syncing).await;
         }
 
+        // Create a forked block that competes with the head block. Both the
+        // head block and this fork block will share the same parent.
+        //
+        // The fork block and head block will both have an unrealized justified
+        // checkpoint at epoch `N` whilst their parent is at `N - 1`.
+        let head_slot = rig.cached_head().head_slot();
+        let parent_slot = head_slot - 1;
+        let fork_block_slot = head_slot + 1;
+        let parent_state = rig
+            .harness
+            .chain
+            .state_at_slot(parent_slot, StateSkipConfig::WithStateRoots)
+            .unwrap();
+        let (fork_block_tuple, _) = rig.harness.make_block(parent_state, fork_block_slot).await;
+        let fork_block = Arc::new(fork_block_tuple.0);
+
         let invalid_head = rig.cached_head();
-        assert_eq!(
-            invalid_head.head_slot() % slots_per_epoch,
-            slots_per_epoch - 1
-        );
 
-        // Advance clock to new epoch to realize the justification of soon-to-be-invalid head block.
-        rig.harness.set_current_slot(invalid_head.head_slot() + 1);
+        // Advance the chain forward two epochs past the current head block.
+        //
+        // This ensures that `voting_source.epoch + 2 >= current_epoch` is
+        // `false` in the `node_is_viable_for_head` function. In effect, this
+        // ensures that no other block but the current head block is viable as a
+        // head block.
+        let invalid_head_epoch = invalid_head.head_slot().epoch(slots_per_epoch);
+        let new_wall_clock_epoch = invalid_head_epoch + 2;
+        rig.harness
+            .set_current_slot(new_wall_clock_epoch.start_slot(slots_per_epoch));
 
         // Invalidate the head block.
         rig.invalidate_manually(invalid_head.head_block_root())
             .await;
 
+        // Since our setup ensures that there is only a single, invalid block
+        // that's viable to for head (according to FFG filtering), setting the
+        // head block as invalid should not result in another head being chosen.
+        // Rather, it should fail to run fork choice and leave invalid block as
+        // the head.
         assert!(rig
             .canonical_head()
             .head_execution_status()
             .unwrap()
             .is_invalid());
 
-        // Finding a new head should fail since the only possible head is not valid.
+        // Ensure that we're getting the correct error when trying to find a new
+        // head.
         rig.assert_get_head_error_contains("InvalidBestNode");
 
         Self {
             rig,
-            fork_block: opt_fork_block.unwrap(),
+            fork_block,
             invalid_head,
         }
     }