diff --git a/tests/core/pyspec/eth2spec/test/helpers/attestations.py b/tests/core/pyspec/eth2spec/test/helpers/attestations.py
index 1e0560405..2bfe63bd1 100644
--- a/tests/core/pyspec/eth2spec/test/helpers/attestations.py
+++ b/tests/core/pyspec/eth2spec/test/helpers/attestations.py
@@ -90,7 +90,7 @@ def build_attestation_data(spec, state, slot, index, shard_transition=None, on_t
                 attestation_data.shard_head_root = shard_transition.shard_data_roots[lastest_shard_data_root_index]
                 attestation_data.shard_transition_root = shard_transition.hash_tree_root()
             else:
-                attestation_data.shard_head_root = state.shard_states[shard].transition_digest
+                attestation_data.shard_head_root = state.shard_states[shard].latest_block_root
                 attestation_data.shard_transition_root = spec.Root()
     return attestation_data
 
diff --git a/tests/core/pyspec/eth2spec/test/helpers/shard_block.py b/tests/core/pyspec/eth2spec/test/helpers/shard_block.py
index 15443e386..ddf66f6c2 100644
--- a/tests/core/pyspec/eth2spec/test/helpers/shard_block.py
+++ b/tests/core/pyspec/eth2spec/test/helpers/shard_block.py
@@ -62,8 +62,6 @@ def get_shard_transitions(spec, parent_beacon_state, shard_blocks):
             on_time_slot,
         )
         len_offset_slots = len(offset_slots)
-        # TODO this is actually unsafe for long offset_slots
-        assert len_offset_slots == on_time_slot - parent_beacon_state.shard_states[shard].slot - 1
         shard_transition = spec.get_shard_transition(parent_beacon_state, shard, blocks)
 
         if len(blocks) > 0: