Add activate and partial withdrawal tests

tests/core/pyspec/eth2spec/test/capella/sanity/test_blocks.py

@@ -1,17 +1,25 @@
+from eth2spec.test.helpers.constants import MINIMAL
 from eth2spec.test.context import (
-    with_capella_and_later, spec_state_test
+    with_capella_and_later,
+    spec_state_test,
+    with_presets,
 )
+from eth2spec.test.helpers.keys import pubkeys
 from eth2spec.test.helpers.state import (
-    state_transition_and_sign_block,
     next_epoch_via_block,
+    state_transition_and_sign_block,
+    transition_to,
+    transition_to_slot_via_block,
+    next_epoch,
+    next_slot,
 )
 from eth2spec.test.helpers.block import (
     build_empty_block_for_next_slot,
     build_empty_block,
 )
 from eth2spec.test.helpers.bls_to_execution_changes import get_signed_address_change
-from eth2spec.test.helpers.state import (
-    next_slot,
+from eth2spec.test.helpers.attestations import (
+    next_epoch_with_attestations,
 )
 from eth2spec.test.helpers.withdrawals import (
     set_eth1_withdrawal_credential_with_balance,
@@ -333,3 +341,71 @@ def test_top_up_and_fully_withdrawal_validator(spec, state):
     balance = state.balances[validator_index]
     current_epoch = spec.get_current_epoch(state)
     assert spec.is_fully_withdrawable_validator(validator, balance, current_epoch)
+
+
+def _insert_validator(spec, state, balance):
+    effective_balance = balance if balance < spec.MAX_EFFECTIVE_BALANCE else spec.MAX_EFFECTIVE_BALANCE
+    validator_index = len(state.validators)
+    validator = spec.Validator(
+        pubkey=pubkeys[validator_index],
+        withdrawal_credentials=spec.ETH1_ADDRESS_WITHDRAWAL_PREFIX + b'\x00' * 11 + b'\x56' * 20,
+        activation_eligibility_epoch=1,
+        activation_epoch=2,
+        exit_epoch=spec.FAR_FUTURE_EPOCH,
+        withdrawable_epoch=spec.FAR_FUTURE_EPOCH,
+        effective_balance=effective_balance,
+    )
+    state.validators.append(validator)
+    state.balances.append(balance)
+    state.previous_epoch_participation.append(spec.ParticipationFlags(0b0000_0000))
+    state.current_epoch_participation.append(spec.ParticipationFlags(0b0000_0000))
+    state.inactivity_scores.append(0)
+
+    return validator_index
+
+
+def _run_activate_and_partial_withdrawal(spec, state, initial_balance):
+    validator_index = _insert_validator(spec, state, balance=initial_balance)
+
+    next_epoch(spec, state)
+    transition_to(spec, state, spec.compute_start_slot_at_epoch(2) - 1)
+
+    assert not spec.is_active_validator(state.validators[validator_index], spec.get_current_epoch(state))
+
+    yield 'pre', state
+
+    blocks = []
+    # To activate
+    signed_block = transition_to_slot_via_block(spec, state, state.slot + 1)
+    blocks.append(signed_block)
+
+    assert spec.is_active_validator(state.validators[validator_index], spec.get_current_epoch(state))
+
+    if initial_balance > spec.MAX_EFFECTIVE_BALANCE:
+        assert spec.is_partially_withdrawable_validator(
+            state.validators[validator_index], state.balances[validator_index])
+    else:
+        assert not spec.is_partially_withdrawable_validator(
+            state.validators[validator_index], state.balances[validator_index])
+
+    # Getting attester rewards and getting partial withdrawals
+    for _ in range(2):
+        _, new_blocks, state = next_epoch_with_attestations(spec, state, True, True)
+        blocks += new_blocks
+
+    yield 'blocks', blocks
+    yield 'post', state
+
+
+@with_capella_and_later
+@with_presets([MINIMAL], reason="too many validators with mainnet config")
+@spec_state_test
+def test_activate_and_partial_withdrawal_max_effective_balance(spec, state):
+    yield from _run_activate_and_partial_withdrawal(spec, state, initial_balance=spec.MAX_EFFECTIVE_BALANCE)
+
+
+@with_capella_and_later
+@with_presets([MINIMAL], reason="too many validators with mainnet config")
+@spec_state_test
+def test_activate_and_partial_withdrawal_overdeposit(spec, state):
+    yield from _run_activate_and_partial_withdrawal(spec, state, initial_balance=spec.MAX_EFFECTIVE_BALANCE + 10000000)

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

@@ -38,8 +38,9 @@ def transition_to_slot_via_block(spec, state, slot):
     Transition to ``slot`` via an empty block transition
     """
     assert state.slot < slot
-    apply_empty_block(spec, state, slot)
+    signed_block = apply_empty_block(spec, state, slot)
     assert state.slot == slot
+    return signed_block
 
 
 def transition_to_valid_shard_slot(spec, state):