Polish withdrawal tests

tests/core/pyspec/eth2spec/test/capella/block_processing/test_process_withdrawals.py

@@ -2,7 +2,6 @@
 
 from eth2spec.test.context import (
     spec_state_test,
-    expect_assertion_error,
     with_presets,
     with_capella_and_later,
 )

tests/core/pyspec/eth2spec/test/electra/block_processing/test_process_withdrawals.py

@@ -2,30 +2,18 @@
 
 from eth2spec.test.context import (
     spec_state_test,
-    with_presets,
     with_electra_and_later,
 )
-from eth2spec.test.helpers.constants import MAINNET, MINIMAL
 from eth2spec.test.helpers.execution_payload import (
     build_empty_execution_payload,
-    compute_el_block_hash,
-)
-from eth2spec.test.helpers.random import (
-    randomize_state,
 )
 from eth2spec.test.helpers.state import (
-    next_epoch,
     next_slot,
 )
 from eth2spec.test.helpers.withdrawals import (
-    get_expected_withdrawals,
-    prepare_expected_withdrawals,
-    set_eth1_withdrawal_credential_with_balance,
-    set_validator_fully_withdrawable,
-    set_validator_partially_withdrawable,
     prepare_expected_withdrawals_compounding,
     run_withdrawals_processing,
-    set_compounding_withdrawal_credential,
+    set_compounding_withdrawal_credential_with_balance,
     prepare_pending_withdrawal,
 )
 
@@ -56,11 +44,10 @@ def test_success_mixed_fully_and_partial_withdrawable_compounding(spec, state):
 def test_success_no_max_effective_balance_compounding(spec, state):
     validator_index = len(state.validators) // 2
     # To be partially withdrawable, the validator's effective balance must be maxed out
-    set_compounding_withdrawal_credential(spec, state, validator_index)
-    validator = state.validators[validator_index]
-    validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA - spec.EFFECTIVE_BALANCE_INCREMENT
-    state.balances[validator_index] = validator.effective_balance
+    effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA - spec.EFFECTIVE_BALANCE_INCREMENT
+    set_compounding_withdrawal_credential_with_balance(spec, state, validator_index, effective_balance)
 
+    validator = state.validators[validator_index]
     assert not spec.is_partially_withdrawable_validator(validator, state.balances[validator_index])
 
     execution_payload = build_empty_execution_payload(spec, state)
@@ -73,11 +60,9 @@ def test_success_no_max_effective_balance_compounding(spec, state):
 def test_success_no_excess_balance_compounding(spec, state):
     validator_index = len(state.validators) // 2
     # To be partially withdrawable, the validator needs an excess balance
-    set_compounding_withdrawal_credential(spec, state, validator_index)
-    validator = state.validators[validator_index]
-    validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
-    state.balances[validator_index] = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
+    set_compounding_withdrawal_credential_with_balance(spec, state, validator_index, spec.MAX_EFFECTIVE_BALANCE_ELECTRA)
 
+    validator = state.validators[validator_index]
     assert not spec.is_partially_withdrawable_validator(validator, state.balances[validator_index])
 
     execution_payload = build_empty_execution_payload(spec, state)
@@ -90,11 +75,11 @@ def test_success_no_excess_balance_compounding(spec, state):
 def test_success_excess_balance_but_no_max_effective_balance_compounding(spec, state):
     validator_index = len(state.validators) // 2
     # To be partially withdrawable, the validator needs both a maxed out effective balance and an excess balance
-    set_compounding_withdrawal_credential(spec, state, validator_index)
-    validator = state.validators[validator_index]
-    validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA - spec.EFFECTIVE_BALANCE_INCREMENT
-    state.balances[validator_index] = spec.MAX_EFFECTIVE_BALANCE_ELECTRA + spec.EFFECTIVE_BALANCE_INCREMENT
+    effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA - spec.EFFECTIVE_BALANCE_INCREMENT
+    balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA + spec.EFFECTIVE_BALANCE_INCREMENT
+    set_compounding_withdrawal_credential_with_balance(spec, state, validator_index, effective_balance, balance)
 
+    validator = state.validators[validator_index]
     assert not spec.is_partially_withdrawable_validator(validator, state.balances[validator_index])
 
     execution_payload = build_empty_execution_payload(spec, state)
@@ -113,9 +98,9 @@ def test_pending_withdrawals_one_skipped_one_effective(spec, state):
 
     # If validator doesn't have an excess balance pending withdrawal is skipped
     state.balances[index_0] = spec.MIN_ACTIVATION_BALANCE
-    
+
     execution_payload = build_empty_execution_payload(spec, state)
-
+    assert state.pending_partial_withdrawals == [withdrawal_0, withdrawal_1]
     yield from run_withdrawals_processing(spec, state, execution_payload, num_expected_withdrawals=1)
 
     assert state.pending_partial_withdrawals == []

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

@@ -84,6 +84,19 @@ def set_compounding_withdrawal_credential(spec, state, index, address=None):
     validator.withdrawal_credentials = spec.COMPOUNDING_WITHDRAWAL_PREFIX + b'\x00' * 11 + address
 
 
+def set_compounding_withdrawal_credential_with_balance(spec, state, index,
+                                                       effective_balance=None, balance=None, address=None):
+    set_compounding_withdrawal_credential(spec, state, index, address)
+
+    if effective_balance is None:
+        effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
+    if balance is None:
+        balance = effective_balance
+
+    state.validators[index].effective_balance = effective_balance
+    state.balances[index] = balance
+
+
 def prepare_expected_withdrawals_compounding(spec, state, rng,
                                              num_full_withdrawals=0,
                                              num_partial_withdrawals_sweep=0,
@@ -95,11 +108,8 @@ def prepare_expected_withdrawals_compounding(spec, state, rng,
     )
 
     for index in fully_withdrawable_indices + partial_withdrawals_sweep_indices:
-        validator = state.validators[index]
-        validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
-        state.balances[index] = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
         address = state.validators[index].withdrawal_credentials[12:]
-        set_compounding_withdrawal_credential(spec, state, index, address=address)
+        set_compounding_withdrawal_credential_with_balance(spec, state, index, address=address)
 
     for index in fully_withdrawable_indices:
         set_validator_fully_withdrawable(spec, state, index)
@@ -113,9 +123,10 @@ def prepare_pending_withdrawal(spec, state, validator_index,
                                effective_balance=32_000_000_000, amount=1_000_000_000):
     assert is_post_electra(spec)
 
-    set_compounding_withdrawal_credential(spec, state, validator_index)
-    state.validators[validator_index].effective_balance = effective_balance
-    state.balances[validator_index] = effective_balance + amount
+    balance = effective_balance + amount
+    set_compounding_withdrawal_credential_with_balance(
+        spec, state, validator_index, effective_balance, balance
+    )
 
     withdrawal = spec.PendingPartialWithdrawal(
         index=validator_index,
@@ -130,6 +141,7 @@ def prepare_pending_withdrawal(spec, state, validator_index,
 # Run processing
 #
 
+
 def verify_post_state(state, spec, expected_withdrawals,
                       fully_withdrawable_indices, partial_withdrawals_indices):
     # Consider verifying also the condition when no withdrawals are expected.