Fix get_pow_block_at_terminal_total_difficulty and add unit tests

specs/merge/validator.md

@@ -43,11 +43,17 @@ Please see related Beacon Chain doc before continuing and use them as a referenc
 ```python
 def get_pow_block_at_terminal_total_difficulty(pow_chain: Dict[Hash32, PowBlock]) -> Optional[PowBlock]:
     # `pow_chain` abstractly represents all blocks in the PoW chain
-    for block in pow_chain:
+    for block in pow_chain.values():
         block_reached_ttd = block.total_difficulty >= TERMINAL_TOTAL_DIFFICULTY
-        # If genesis block, no parent exists so reaching TTD alone qualifies as valid terminal block
-        if block_reached_ttd and block.parent_hash == Hash32():
-            return block
+        # Genesis block
+        if block.parent_hash == Hash32():
+            # No parent exists, so reaching TTD alone qualifies as valid terminal block
+            if block_reached_ttd:
+                return block
+            # Continue to iterate the next block
+            if block.parent_hash not in pow_chain:
+                continue
+
         parent = pow_chain[block.parent_hash]
         parent_reached_ttd = parent.total_difficulty >= TERMINAL_TOTAL_DIFFICULTY
         if block_reached_ttd and not parent_reached_ttd:

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

@@ -15,6 +15,12 @@ def head(self, offset=0):
         assert offset <= 0
         return self.blocks[offset - 1]
 
+    def to_dict(self):
+        return {
+            block.block_hash: block
+            for block in self.blocks
+        }
+
 
 def prepare_random_pow_block(spec, rng=Random(3131)):
     return spec.PowBlock(

tests/core/pyspec/eth2spec/test/merge/unittests/validator/test_validator.py

@@ -0,0 +1,64 @@
+from eth2spec.test.helpers.pow_block import (
+    prepare_random_pow_chain,
+)
+from eth2spec.test.context import (
+    spec_state_test,
+    with_merge_and_later,
+)
+
+
+# For test_get_pow_block_at_terminal_total_difficulty
+IS_HEAD_BLOCK = 'is_head_block'
+IS_HEAD_PARENT_BLOCK = 'is_head_parent_block'
+
+# NOTE: The following parameter names are in the view of the head block (the second block)
+# 'block_reached_ttd', 'block_parent_hash_is_empty', 'parent_reached_ttd', 'return_block'
+expected_results = [
+    (False, False, False, None),
+    (False, False, True, IS_HEAD_PARENT_BLOCK),
+    (False, True, False, None),
+    (False, True, True, IS_HEAD_PARENT_BLOCK),
+    (True, False, False, IS_HEAD_BLOCK),
+    (True, False, True, IS_HEAD_PARENT_BLOCK),
+    (True, True, False, IS_HEAD_BLOCK),
+    (True, True, True, IS_HEAD_PARENT_BLOCK),
+]
+# NOTE: since the first block's `parent_hash` is set to `Hash32()` in test, if `parent_reached_ttd is True`,
+# it would return the first block (IS_HEAD_PARENT_BLOCK).
+
+
+@with_merge_and_later
+@spec_state_test
+def test_get_pow_block_at_terminal_total_difficulty(spec, state):
+    for result in expected_results:
+        (
+            block_reached_ttd,
+            block_parent_hash_is_empty,
+            parent_reached_ttd,
+            return_block
+        ) = result
+        pow_chain = prepare_random_pow_chain(spec, 2)
+        pow_chain.head(-1).parent_hash = spec.Hash32()
+
+        if block_reached_ttd:
+            pow_chain.head().total_difficulty = spec.config.TERMINAL_TOTAL_DIFFICULTY
+        else:
+            pow_chain.head().total_difficulty = spec.config.TERMINAL_TOTAL_DIFFICULTY - 1
+
+        if parent_reached_ttd:
+            pow_chain.head(-1).total_difficulty = spec.config.TERMINAL_TOTAL_DIFFICULTY
+        else:
+            pow_chain.head(-1).total_difficulty = spec.config.TERMINAL_TOTAL_DIFFICULTY - 1
+
+        if block_parent_hash_is_empty:
+            pow_chain.head().parent_hash = spec.Hash32()
+
+        pow_block = spec.get_pow_block_at_terminal_total_difficulty(pow_chain.to_dict())
+        if return_block == IS_HEAD_BLOCK:
+            assert pow_block == pow_chain.head()
+        elif return_block == IS_HEAD_PARENT_BLOCK:
+            assert pow_block == pow_chain.head(-1)
+        elif return_block is None:
+            assert pow_block is None
+        else:
+            raise Exception('Something is wrong')