testing: rearrange keys file

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

@@ -44,7 +44,6 @@
     LEAN_ENV_TO_SCHEMES,
     XmssKeyManager,
     create_dummy_signature,
-    get_shared_key_manager,
 )
 from ..test_types import (
     AggregatedAttestationSpec,
@@ -191,7 +190,7 @@ def make_fixture(self) -> Self:
         # XMSS keys are expensive to generate.
         # The shared key manager caches keys across tests.
         # Tests requiring higher max slot trigger key expansion.
-        key_manager = get_shared_key_manager(max_slot=self.max_slot)
+        key_manager = XmssKeyManager.shared(max_slot=self.max_slot)
 
         # Validator pubkey synchronization
         #

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

@@ -27,7 +27,7 @@
 from lean_spec.subspecs.xmss.aggregation import AggregatedSignatureProof
 from lean_spec.types import ByteListMiB, Bytes32
 
-from ..keys import XmssKeyManager, create_dummy_signature, get_shared_key_manager
+from ..keys import XmssKeyManager, create_dummy_signature
 from ..test_types import AggregatedAttestationSpec, BlockSpec
 from .base import BaseConsensusFixture
 
@@ -122,7 +122,7 @@ def make_fixture(self) -> VerifySignaturesTest:
         assert self.anchor_state is not None, "anchor state must be set before making the fixture"
 
         # Use shared key manager
-        key_manager = get_shared_key_manager()
+        key_manager = XmssKeyManager.shared()
 
         # Build the signed block
         signed_block = self._build_block_from_spec(self.block, self.anchor_state, key_manager)

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

@@ -6,7 +6,7 @@
 from lean_spec.subspecs.containers.validator import Validator, ValidatorIndex
 from lean_spec.types import Bytes52, Uint64
 
-from ..keys import get_shared_key_manager
+from ..keys import XmssKeyManager
 
 
 def generate_pre_state(**kwargs: Any) -> State:
@@ -24,7 +24,7 @@ def generate_pre_state(**kwargs: Any) -> State:
     genesis_time = kwargs.get("genesis_time", Uint64(0))
     num_validators = kwargs.get("num_validators", 4)
 
-    key_manager = get_shared_key_manager()
+    key_manager = XmssKeyManager.shared()
     available_keys = len(key_manager)
 
     assert num_validators <= available_keys, (

tests/consensus/devnet/ssz/test_xmss_containers.py

@@ -2,7 +2,7 @@
 
 import pytest
 from consensus_testing import SSZTestFiller
-from consensus_testing.keys import create_dummy_signature, get_shared_key_manager
+from consensus_testing.keys import XmssKeyManager, create_dummy_signature
 
 from lean_spec.subspecs.containers import ValidatorIndex
 from lean_spec.subspecs.containers.attestation import AggregationBits
@@ -52,9 +52,9 @@ def test_signature_zero(ssz: SSZTestFiller) -> None:
 
 def test_signature_actual(ssz: SSZTestFiller) -> None:
     """SSZ roundtrip for a cryptographically valid Signature produced by signing."""
-    key_manager = get_shared_key_manager()
+    key_manager = XmssKeyManager.shared()
     scheme = key_manager.scheme
-    sk = key_manager.keys[ValidatorIndex(0)].attestation_secret
+    sk = key_manager[ValidatorIndex(0)].attestation_secret
     signature = scheme.sign(sk, Slot(0), Bytes32(b"\x42" * 32))
     ssz(type_name="Signature", value=signature)
 

tests/lean_spec/conftest.py

@@ -10,7 +10,7 @@
 from collections.abc import Callable
 
 import pytest
-from consensus_testing.keys import XmssKeyManager, get_shared_key_manager
+from consensus_testing.keys import XmssKeyManager
 
 from lean_spec.subspecs.containers import Block, State
 from lean_spec.subspecs.containers.slot import Slot
@@ -28,7 +28,7 @@
 @pytest.fixture
 def key_manager() -> XmssKeyManager:
     """XMSS key manager for signing attestations."""
-    return get_shared_key_manager(max_slot=Slot(20))
+    return XmssKeyManager.shared(max_slot=Slot(20))
 
 
 _DEFAULT_VALIDATOR_ID = ValidatorIndex(0)

tests/lean_spec/helpers/builders.py

@@ -9,7 +9,7 @@
 
 from typing import NamedTuple, cast
 
-from consensus_testing.keys import XmssKeyManager, get_shared_key_manager
+from consensus_testing.keys import XmssKeyManager
 
 from lean_spec.subspecs.chain.clock import Interval, SlotClock
 from lean_spec.subspecs.chain.config import INTERVALS_PER_SLOT
@@ -101,20 +101,18 @@ def make_validators(count: int) -> Validators:
 
 def make_validators_from_key_manager(key_manager: XmssKeyManager, count: int) -> Validators:
     """Build a validator registry with real XMSS keys from a key manager."""
-    return Validators(
-        data=[
+    validators = []
+    for i in range(count):
+        idx = ValidatorIndex(i)
+        att_pk, prop_pk = key_manager.get_public_keys(idx)
+        validators.append(
             Validator(
-                attestation_pubkey=Bytes52(
-                    key_manager[ValidatorIndex(i)].attestation_public.encode_bytes()
-                ),
-                proposal_pubkey=Bytes52(
-                    key_manager[ValidatorIndex(i)].proposal_public.encode_bytes()
-                ),
-                index=ValidatorIndex(i),
+                attestation_pubkey=Bytes52(att_pk.encode_bytes()),
+                proposal_pubkey=Bytes52(prop_pk.encode_bytes()),
+                index=idx,
             )
-            for i in range(count)
-        ]
-    )
+        )
+    return Validators(data=validators)
 
 
 def make_genesis_state(
@@ -413,7 +411,7 @@ def make_keyed_genesis_state(
 ) -> State:
     """Create a genesis state with real XMSS keys from the shared key manager."""
     if key_manager is None:
-        key_manager = get_shared_key_manager()
+        key_manager = XmssKeyManager.shared()
     validators = make_validators_from_key_manager(key_manager, num_validators)
     return make_genesis_state(validators=validators)
 
@@ -428,7 +426,7 @@ def make_aggregated_proof(
     xmss_participants = AggregationBits.from_validator_indices(ValidatorIndices(data=participants))
     raw_xmss = list(
         zip(
-            [key_manager[vid].attestation_public for vid in participants],
+            [key_manager.get_public_keys(vid)[0] for vid in participants],
             [key_manager.sign_attestation_data(vid, attestation_data) for vid in participants],
             strict=True,
         )

tests/lean_spec/subspecs/containers/conftest.py

@@ -3,12 +3,12 @@
 from __future__ import annotations
 
 import pytest
-from consensus_testing.keys import XmssKeyManager, get_shared_key_manager
+from consensus_testing.keys import XmssKeyManager
 
 from lean_spec.subspecs.containers.slot import Slot
 
 
 @pytest.fixture
 def container_key_manager() -> XmssKeyManager:
     """Key manager for container tests."""
-    return get_shared_key_manager(max_slot=Slot(20))
+    return XmssKeyManager.shared(max_slot=Slot(20))

tests/lean_spec/subspecs/validator/test_service.py

@@ -5,7 +5,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
-from consensus_testing.keys import XmssKeyManager, get_shared_key_manager
+from consensus_testing.keys import XmssKeyManager
 
 from lean_spec.subspecs.chain.clock import SlotClock
 from lean_spec.subspecs.chain.config import MILLISECONDS_PER_INTERVAL
@@ -445,7 +445,7 @@ class TestValidatorServiceIntegration:
     @pytest.fixture
     def key_manager(self) -> XmssKeyManager:
         """Key manager with pre-generated test keys."""
-        return get_shared_key_manager(max_slot=Slot(10))
+        return XmssKeyManager.shared(max_slot=Slot(10))
 
     @pytest.fixture
     def real_store(self, key_manager: XmssKeyManager) -> Store: