diff --git a/tests/core/pyspec/eth2spec/test/phase0/unittests/fork_choice/test_on_attestation.py b/tests/core/pyspec/eth2spec/test/phase0/unittests/fork_choice/test_on_attestation.py
index 9007a778f..b05f56ba9 100644
--- a/tests/core/pyspec/eth2spec/test/phase0/unittests/fork_choice/test_on_attestation.py
+++ b/tests/core/pyspec/eth2spec/test/phase0/unittests/fork_choice/test_on_attestation.py
@@ -1,7 +1,7 @@
 from eth2spec.test.context import with_all_phases, spec_state_test
 from eth2spec.test.helpers.block import build_empty_block_for_next_slot
 from eth2spec.test.helpers.attestations import get_valid_attestation, sign_attestation
-from eth2spec.test.helpers.constants import PHASE0, ALTAIR, MERGE
+from eth2spec.test.helpers.constants import PHASE0, ALTAIR, MERGE, CAPELLA
 from eth2spec.test.helpers.state import transition_to, state_transition_and_sign_block, next_epoch, next_slot
 from eth2spec.test.helpers.fork_choice import get_genesis_forkchoice_store
 
@@ -19,7 +19,7 @@ def run_on_attestation(spec, state, store, attestation, valid=True):
     spec.on_attestation(store, attestation)
 
     sample_index = indexed_attestation.attesting_indices[0]
-    if spec.fork in (PHASE0, ALTAIR, MERGE):
+    if spec.fork in (PHASE0, ALTAIR, MERGE, CAPELLA):
         latest_message = spec.LatestMessage(
             epoch=attestation.data.target.epoch,
             root=attestation.data.beacon_block_root,