tests: add test_fork_choice_head test vectors (leanEthereum#103)

* tests: add test_fork_choice_head test vectors

* Update tests/consensus/devnet/fc/test_fork_choice_head.py

Co-authored-by: Unnawut Leepaisalsuwanna <921194+unnawut@users.noreply.github.com>

* fix comment

---------

Co-authored-by: Unnawut Leepaisalsuwanna <921194+unnawut@users.noreply.github.com>

Author: tcoratger

packages/testing/src/consensus_testing/test_fixtures/fork_choice.py

@@ -121,6 +121,12 @@ def make_fixture(self) -> "ForkChoiceTest":
             anchor_block=self.anchor_block,
         )
 
+        # Block registry for label-based fork creation
+        self._block_registry: dict[str, Block] = {}
+
+        # Register genesis/anchor block with implicit label
+        self._block_registry["genesis"] = self.anchor_block
+
         # Process each step
         for i, step in enumerate(self.steps):
             try:
@@ -130,12 +136,23 @@ def make_fixture(self) -> "ForkChoiceTest":
 
                 elif isinstance(step, BlockStep):
                     # Build SignedBlockWithAttestation from BlockSpec
-                    signed_block = self._build_block_from_spec(step.block, store)
+                    signed_block = self._build_block_from_spec(
+                        step.block, store, self._block_registry
+                    )
 
                     # Store the filled Block for serialization
                     block = signed_block.message.block
                     step._filled_block = block
 
+                    # Register block if it has a label
+                    if step.block.label is not None:
+                        if step.block.label in self._block_registry:
+                            raise ValueError(
+                                f"Step {i}: duplicate label '{step.block.label}' - "
+                                f"labels must be unique within a test"
+                            )
+                        self._block_registry[step.block.label] = block
+
                     # Automatically advance time to block's slot before processing
                     # Compute the time corresponding to the block's slot
                     block_time = store.config.genesis_time + block.slot * Uint64(SECONDS_PER_SLOT)
@@ -155,7 +172,9 @@ def make_fixture(self) -> "ForkChoiceTest":
 
                 # Validate checks if provided
                 if step.checks is not None:
-                    step.checks.validate_against_store(store, step_index=i)
+                    step.checks.validate_against_store(
+                        store, step_index=i, block_registry=self._block_registry
+                    )
 
             except Exception as e:
                 if step.valid:
@@ -175,7 +194,12 @@ def make_fixture(self) -> "ForkChoiceTest":
         # Return self (fixture is already complete)
         return self
 
-    def _build_block_from_spec(self, spec: BlockSpec, store: Store) -> SignedBlockWithAttestation:
+    def _build_block_from_spec(
+        self,
+        spec: BlockSpec,
+        store: Store,
+        block_registry: dict[str, Block],
+    ) -> SignedBlockWithAttestation:
         """
         Build a full SignedBlockWithAttestation from a lightweight BlockSpec.
 
@@ -194,6 +218,8 @@ def _build_block_from_spec(self, spec: BlockSpec, store: Store) -> SignedBlockWi
             The lightweight block specification.
         store : Store
             The fork choice store (used to get head state and latest justified).
+        block_registry : dict[str, Block]
+            Registry of labeled blocks for fork creation.
 
         Returns:
         -------
@@ -207,12 +233,31 @@ def _build_block_from_spec(self, spec: BlockSpec, store: Store) -> SignedBlockWi
         else:
             proposer_index = spec.proposer_index
 
-        # Get the current head state from the store
-        head_state = store.states[store.head]
+        # Resolve parent block if parent_label is specified
+        if spec.parent_label is not None:
+            if spec.parent_label not in block_registry:
+                raise ValueError(
+                    f"parent_label '{spec.parent_label}' not found - "
+                    f"available labels: {list(block_registry.keys())}"
+                )
+            parent_block = block_registry[spec.parent_label]
+            parent_root = hash_tree_root(parent_block)
+
+            # Get state at the parent block
+            if parent_root not in store.states:
+                raise ValueError(
+                    f"parent_label '{spec.parent_label}' (root=0x{parent_root.hex()[:16]}...) "
+                    f"has no state in store - cannot build on this fork"
+                )
+            parent_state = store.states[parent_root]
 
-        # Dry-run to build block with correct state root
-        temp_state = head_state.process_slots(spec.slot)
-        parent_root = hash_tree_root(temp_state.latest_block_header)
+            # Advance state to the new block's slot
+            temp_state = parent_state.process_slots(spec.slot)
+        else:
+            # Default: build on current head
+            head_state = store.states[store.head]
+            temp_state = head_state.process_slots(spec.slot)
+            parent_root = hash_tree_root(temp_state.latest_block_header)
 
         # Build body (empty for now, attestations can be added later if needed)
         body = BlockBody(attestations=Attestations(data=[]))

packages/testing/src/consensus_testing/test_types/block_spec.py

@@ -52,3 +52,31 @@ class BlockSpec(CamelModel):
     If None, framework creates empty body for state transition tests,
     or collects attestations for fork choice tests.
     """
+
+    label: str | None = None
+    """
+    Optional label to tag this block for later reference.
+
+    Allows creating forks by building blocks on labeled ancestors:
+    ```python
+    BlockSpec(slot=Slot(1), label="fork_a")  # Tag this block
+    BlockSpec(slot=Slot(2), parent_label="fork_a")  # Build on it
+    ```
+
+    Labels must be unique within a test.
+    """
+
+    parent_label: str | None = None
+    """
+    Optional label referencing a previously created block as parent.
+
+    Enables explicit fork creation:
+    ```python
+    BlockSpec(slot=Slot(1), label="common")
+    BlockSpec(slot=Slot(2), parent_label="common", label="fork_a")
+    BlockSpec(slot=Slot(2), parent_label="common", label="fork_b")
+    ```
+
+    If None, parent is determined by the current canonical head.
+    If specified, parent_root is computed from the labeled block.
+    """

packages/testing/src/consensus_testing/test_types/store_checks.py

@@ -7,6 +7,7 @@
 
 if TYPE_CHECKING:
     from lean_spec.subspecs.containers import SignedAttestation
+    from lean_spec.subspecs.containers.block.block import Block
     from lean_spec.subspecs.forkchoice.store import Store
 
 
@@ -109,6 +110,18 @@ class StoreChecks(CamelModel):
     head_root: Bytes32 | None = None
     """Expected head block root."""
 
+    head_root_label: str | None = None
+    """
+    Expected head block root by label reference.
+
+    Alternative to head_root that uses the block label system.
+    The framework will resolve this label to the actual block root
+    and validate the head matches.
+
+    Example:
+        StoreChecks(head_root_label="fork_a")  # Validates head is fork_a block
+    """
+
     latest_justified_slot: Slot | None = None
     """Expected latest justified checkpoint slot."""
 
@@ -136,7 +149,9 @@ class StoreChecks(CamelModel):
     attestation_checks: list[AttestationCheck] | None = None
     """Optional list of attestation content checks for specific validators."""
 
-    def validate_against_store(self, store: "Store", step_index: int) -> None:
+    def validate_against_store(
+        self, store: "Store", step_index: int, block_registry: dict[str, "Block"] | None = None
+    ) -> None:
         """
         Validate these checks against actual Store state.
 
@@ -149,6 +164,8 @@ def validate_against_store(self, store: "Store", step_index: int) -> None:
             The fork choice store to validate against.
         step_index : int
             Index of the step being validated (for error messages).
+        block_registry : dict[str, Block] | None
+            Optional registry of labeled blocks for resolving head_root_label.
 
         Raises:
         ------
@@ -183,6 +200,35 @@ def validate_against_store(self, store: "Store", step_index: int) -> None:
                         f"expected 0x{expected_value.hex()}"
                     )
 
+            elif field_name == "head_root_label":
+                # Resolve label to root
+                if block_registry is None:
+                    raise ValueError(
+                        f"Step {step_index}: head_root_label='{expected_value}' specified "
+                        f"but block_registry not provided to validate_against_store()"
+                    )
+
+                if expected_value not in block_registry:
+                    available = list(block_registry.keys())
+                    raise ValueError(
+                        f"Step {step_index}: head_root_label='{expected_value}' not found "
+                        f"in block registry. Available labels: {available}"
+                    )
+
+                # Import hash_tree_root locally to avoid circular import
+                from lean_spec.subspecs.ssz import hash_tree_root
+
+                expected_block = block_registry[expected_value]
+                expected_root = hash_tree_root(expected_block)
+                actual_root = store.head
+
+                if actual_root != expected_root:
+                    raise AssertionError(
+                        f"Step {step_index}: head.root = 0x{actual_root.hex()}, "
+                        f"expected 0x{expected_root.hex()} "
+                        f"(label '{expected_value}')"
+                    )
+
             elif field_name == "latest_justified_slot":
                 actual = store.latest_justified.slot
                 if actual != expected_value: