test: simplify validators initialization (#219)

* test: simplify validators init

* fix: linting

* fix: linting

Author: unnawut

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

@@ -1,10 +1,29 @@
 """Consensus layer pre-state generation."""
 
+from functools import lru_cache
 from typing import Any
 
+from lean_spec.subspecs.containers.slot import Slot
 from lean_spec.subspecs.containers.state import State, Validators
 from lean_spec.subspecs.containers.validator import Validator
-from lean_spec.types import Bytes52, Uint64
+from lean_spec.types import Uint64
+
+from ..keys import XmssKeyManager
+
+
+@lru_cache(maxsize=1)
+def _get_shared_key_manager() -> XmssKeyManager:
+    """
+    Get or create the shared XMSS key manager for reusing keys across tests.
+
+    Uses functools.lru_cache to create a singleton instance that's shared
+    across all test fixture generations within a session. This optimizes
+    performance by reusing keys when possible.
+
+    Returns:
+        Shared XmssKeyManager instance with max_slot=10.
+    """
+    return XmssKeyManager(max_slot=Slot(10))
 
 
 def generate_pre_state(**kwargs: Any) -> State:
@@ -14,20 +33,20 @@ def generate_pre_state(**kwargs: Any) -> State:
     Args:
         **kwargs: Optional keyword arguments:
             - genesis_time: The genesis timestamp (defaults to Uint64(0)).
-            - validators: Validators list (defaults to 4 validators with dummy pubkeys).
+            - num_validators: Number of validators (defaults to 4 validators).
 
     Returns:
         State: A properly initialized consensus state.
     """
     genesis_time = kwargs.get("genesis_time", Uint64(0))
-
-    # If validators not provided, create a default set of 4 validators with dummy pubkeys
-    # TODO: Set an appropriate default here for test fixtures
-    if "validators" not in kwargs:
-        validators = Validators(
-            data=[Validator(pubkey=Bytes52.zero(), index=Uint64(i)) for i in range(4)]
-        )
-    else:
-        validators = kwargs["validators"]
+    num_validators = kwargs.get("num_validators", 4)
+    key_manager = _get_shared_key_manager()
+
+    validators = Validators(
+        data=[
+            Validator(pubkey=key_manager[Uint64(i)].public.encode_bytes(), index=Uint64(i))
+            for i in range(num_validators)
+        ]
+    )
 
     return State.generate_genesis(genesis_time=genesis_time, validators=validators)

tests/consensus/devnet/fc/test_fork_choice_reorgs.py

@@ -361,11 +361,7 @@ def test_reorg_with_slot_gaps(
     conditions where perfect block production is impossible.
     """
     fork_choice_test(
-        anchor_state=generate_pre_state(
-            validators=Validators(
-                data=[Validator(pubkey=Bytes52.zero(), index=Uint64(i)) for i in range(10)]
-            ),
-        ),
+        anchor_state=generate_pre_state(num_validators=10),
         steps=[
             # Base at slot 1
             BlockStep(
@@ -577,11 +573,7 @@ def test_reorg_prevention_heavy_fork_resists_light_competition(
     - Network naturally converges on heaviest fork
     """
     fork_choice_test(
-        anchor_state=generate_pre_state(
-            validators=Validators(
-                data=[Validator(pubkey=Bytes52.zero(), index=Uint64(i)) for i in range(12)]
-            )
-        ),
+        anchor_state=generate_pre_state(num_validators=12),
         steps=[
             # Common base
             BlockStep(
@@ -814,11 +806,7 @@ def test_reorg_on_newly_justified_slot(
     """
     fork_choice_test(
         # Using 9 validators: 3 for Fork A and 6 for Fork B to achieve 2/3rd for Fork B
-        anchor_state=generate_pre_state(
-            validators=Validators(
-                data=[Validator(pubkey=Bytes52.zero(), index=Uint64(i)) for i in range(9)]
-            )
-        ),
+        anchor_state=generate_pre_state(num_validators=9),
         steps=[
             # Common base at slot 1
             BlockStep(

tests/consensus/devnet/state_transition/test_genesis.py

@@ -127,13 +127,8 @@ def test_genesis_custom_validator_set(
     Genesis state should contain exactly 8 validators while
     maintaining all other genesis properties.
     """
-    # Create 8 validators with unique pubkeys
-    validators = Validators(
-        data=[Validator(pubkey=Bytes52(bytes([i] * 52)), index=Uint64(i)) for i in range(8)]
-    )
-
     state_transition_test(
-        pre=generate_pre_state(validators=validators),
+        pre=generate_pre_state(num_validators=8),
         blocks=[],
         post=StateExpectation(
             slot=Slot(0),