clean up some withdrawal logic and add tests

specs/capella/beacon-chain.md

@@ -140,7 +140,8 @@ def is_withdrawable_validator(validator: Validator, epoch: Epoch) -> bool:
     """
     Check if ``validator`` is withdrawable.
     """
-    return validator.withdrawable_epoch <= epoch
+    is_eth1_withdrawal_prefix = validator.withdrawal_credentials[0:1] == ETH1_ADDRESS_WITHDRAWAL_PREFIX
+    return is_eth1_withdrawal_prefix and validator.withdrawable_epoch <= epoch < validator.withdrawn_epoch
 ```
 
 ## Beacon chain state transition function
@@ -172,10 +173,8 @@ def process_epoch(state: BeaconState) -> None:
 def process_withdrawals(state: BeaconState) -> None:
     current_epoch = get_current_epoch(state)
     for index, validator in enumerate(state.validators):
-        balance = state.balances[index]
+        if is_withdrawable_validator(validator, current_epoch):
-        is_balance_nonzero = state.balances[index] == 0
+            # TODO, consider the zero-balance case
-        is_eth1_withdrawal_prefix = validator.withdrawal_credentials[0] == ETH1_ADDRESS_WITHDRAWAL_PREFIX
+            withdraw(state, ValidatorIndex(index), state.balances[index])
-        if is_balance_nonzero and is_eth1_withdrawal_prefix and is_withdrawable_validator(validator, current_epoch):
-            withdraw(state, ValidatorIndex(index), balance)
             validator.withdrawn_epoch = current_epoch
 ```

tests/core/pyspec/eth2spec/test/altair/epoch_processing/test_process_withdrawals.py

@@ -0,0 +1,95 @@
+from random import Random
+
+from eth2spec.test.helpers.constants import MINIMAL
+from eth2spec.test.context import (
+    with_capella_and_later,
+    spec_state_test,
+)
+from eth2spec.test.helpers.epoch_processing import run_epoch_processing_with
+
+
+def set_validator_withdrawable(spec, state, index, withdrawable_epoch=None):
+    if withdrawable_epoch is None:
+        withdrawable_epoch = spec.get_current_epoch(state)
+
+    validator = state.validators[index]
+    validator.withdrawable_epoch = withdrawable_epoch
+    validator.withdrawal_credentials = spec.ETH1_ADDRESS_WITHDRAWAL_PREFIX + validator.withdrawal_credentials[1:]
+
+    assert spec.is_withdrawable_validator(validator, withdrawable_epoch)
+
+
+def run_process_withdrawals(spec, state, num_expected_withdrawals=None):
+    to_be_withdrawn_indices = [
+        index for index, validator in enumerate(state.validators)
+        if spec.is_withdrawable_validator(validator, spec.get_current_epoch(state))
+    ]
+
+    if num_expected_withdrawals is not None:
+        assert len(to_be_withdrawn_indices) == num_expected_withdrawals
+
+    yield from run_epoch_processing_with(spec, state, 'process_withdrawals')
+
+    for index in to_be_withdrawn_indices:
+        validator = state.validators[index]
+        assert validator.withdrawn_epoch == spec.get_current_epoch(state)
+        assert state.balances[index] == 0
+
+
+@with_capella_and_later
+@spec_state_test
+def test_no_withdrawals(spec, state):
+    pre_validators = state.validators.copy()
+    yield from run_process_withdrawals(spec, state, 0)
+
+    assert pre_validators == state.validators
+
+
+@with_capella_and_later
+@spec_state_test
+def test_no_withdrawals_but_some_next_epoch(spec, state):
+    current_epoch = spec.get_current_epoch(state)
+
+    # Make a few validators withdrawable at the *next* epoch
+    for index in range(3):
+        set_validator_withdrawable(spec, state, index, current_epoch + 1)
+
+    yield from run_process_withdrawals(spec, state, 0)
+
+
+@with_capella_and_later
+@spec_state_test
+def test_single_withdrawal(spec, state):
+    current_epoch = spec.get_current_epoch(state)
+
+    # Make one validator withdrawable
+    set_validator_withdrawable(spec, state, current_epoch)
+
+    yield from run_process_withdrawals(spec, state, 1)
+
+
+@with_capella_and_later
+@spec_state_test
+def test_multi_withdrawal(spec, state):
+    current_epoch = spec.get_current_epoch(state)
+
+    # Make a few validators withdrawable
+    for index in range(3):
+        set_validator_withdrawable(spec, state, index)
+
+    yield from run_process_withdrawals(spec, state, 3)
+
+
+@with_capella_and_later
+@spec_state_test
+def test_all_withdrawal(spec, state):
+    current_epoch = spec.get_current_epoch(state)
+
+    # Make all validators withdrawable
+    for index in range(len(state.validators)):
+        set_validator_withdrawable(spec, state, index)
+
+    yield from run_process_withdrawals(spec, state, len(state.validators))
+
+
+

tests/core/pyspec/eth2spec/test/helpers/capella/fork.py

@@ -46,7 +46,7 @@ def run_fork_test(post_spec, pre_state):
             'pubkey', 'withdrawal_credentials',
             'effective_balance',
             'slashed',
-            'activation_eligibility_epoch', 'activation_epoch', 'exit_epoch', 'withdrawable_epoch'
+            'activation_eligibility_epoch', 'activation_epoch', 'exit_epoch', 'withdrawable_epoch',
         ]
         for field in stable_validator_fields:
             assert getattr(pre_validator, field) == getattr(post_validator, field)

tests/core/pyspec/eth2spec/test/helpers/epoch_processing.py

@@ -28,6 +28,7 @@ def get_process_calls(spec):
             'process_participation_record_updates'
         ),
         'process_sync_committee_updates',  # altair
+        'process_withdrawals',  # capella
         # TODO: add sharding processing functions when spec stabilizes.
     ]
 

tests/core/pyspec/eth2spec/test/helpers/genesis.py

@@ -1,6 +1,6 @@
 from eth2spec.test.helpers.constants import (
     ALTAIR, MERGE,
-    FORKS_BEFORE_ALTAIR, FORKS_BEFORE_MERGE,
+    FORKS_BEFORE_ALTAIR, FORKS_BEFORE_MERGE, FORKS_BEFORE_CAPELLA,
 )
 from eth2spec.test.helpers.keys import pubkeys
 
@@ -9,7 +9,7 @@ def build_mock_validator(spec, i: int, balance: int):
     pubkey = pubkeys[i]
     # insecurely use pubkey as withdrawal key as well
     withdrawal_credentials = spec.BLS_WITHDRAWAL_PREFIX + spec.hash(pubkey)[1:]
-    return spec.Validator(
+    validator = spec.Validator(
         pubkey=pubkeys[i],
         withdrawal_credentials=withdrawal_credentials,
         activation_eligibility_epoch=spec.FAR_FUTURE_EPOCH,
@@ -19,6 +19,11 @@ def build_mock_validator(spec, i: int, balance: int):
         effective_balance=min(balance - balance % spec.EFFECTIVE_BALANCE_INCREMENT, spec.MAX_EFFECTIVE_BALANCE)
     )
 
+    if spec.fork not in FORKS_BEFORE_CAPELLA:
+        validator.withdrawn_epoch = spec.FAR_FUTURE_EPOCH
+
+    return validator
+
 
 def get_sample_genesis_execution_payload_header(spec,
                                                 eth1_block_hash=None):