address issue in looking up cached block generator (#17622)

chia/full_node/full_node.py

@@ -1674,15 +1674,18 @@ async def add_block(
             # This is the case where we already had the unfinished block, and asked for this block without
             # the transactions (since we already had them). Therefore, here we add the transactions.
             unfinished_rh: bytes32 = block.reward_chain_block.get_unfinished().get_hash()
-            unf_block: Optional[UnfinishedBlock] = self.full_node_store.get_unfinished_block(unfinished_rh)
+            foliage_hash: Optional[bytes32] = block.foliage.foliage_transaction_block_hash
+            assert foliage_hash is not None
+            unf_block: Optional[UnfinishedBlock]
+            unf_block, count, has_better = self.full_node_store.get_unfinished_block2(unfinished_rh, foliage_hash)
             if (
                 unf_block is not None
                 and unf_block.transactions_generator is not None
                 and unf_block.foliage_transaction_block == block.foliage_transaction_block
             ):
                 # We checked that the transaction block is the same, therefore all transactions and the signature
                 # must be identical in the unfinished and finished blocks. We can therefore use the cache.
-                pre_validation_result = self.full_node_store.get_unfinished_block_result(unfinished_rh)
+                pre_validation_result = self.full_node_store.get_unfinished_block_result2(unfinished_rh, foliage_hash)
                 assert pre_validation_result is not None
                 block = dataclasses.replace(
                     block,

chia/full_node/full_node_store.py

@@ -268,20 +268,12 @@ def get_unfinished_block2(
         else:
             return entry.unfinished_block, len(result), has_better
 
-    def get_unfinished_block_result(self, unfinished_reward_hash: bytes32) -> Optional[PreValidationResult]:
-        result = self.unfinished_blocks.get(unfinished_reward_hash, None)
-        if result is None:
-            return None
-        return next(iter(result.values())).result
-
     def get_unfinished_block_result2(
-        self, unfinished_reward_hash: bytes32, unfinished_foliage_hash: Optional[bytes32]
+        self, unfinished_reward_hash: bytes32, unfinished_foliage_hash: bytes32
     ) -> Optional[PreValidationResult]:
         result = self.unfinished_blocks.get(unfinished_reward_hash, None)
         if result is None:
             return None
-        if unfinished_foliage_hash is None:
-            return next(iter(result.values())).result
         else:
             entry = result.get(unfinished_foliage_hash)
             return None if entry is None else entry.result

tests/core/full_node/stores/test_full_node_store.py

@@ -224,13 +224,11 @@ async def test_basic_store(
             foliage_hash is not None and dummy_hash > foliage_hash,
         )
 
-        ublock = store.get_unfinished_block_result(unf_block.partial_hash)
-        assert ublock is not None and ublock.required_iters == uint64(123532)
-        ublock = store.get_unfinished_block_result2(
-            unf_block.partial_hash, unf_block.foliage.foliage_transaction_block_hash
-        )
-
-        assert ublock is not None and ublock.required_iters == uint64(123532)
+        if unf_block.foliage.foliage_transaction_block_hash:
+            ublock = store.get_unfinished_block_result2(
+                unf_block.partial_hash, unf_block.foliage.foliage_transaction_block_hash
+            )
+            assert ublock is not None and ublock.required_iters == uint64(123532)
 
         store.remove_unfinished_block(unf_block.partial_hash)
         assert store.get_unfinished_block(unf_block.partial_hash) is None
@@ -289,17 +287,11 @@ async def test_basic_store(
     )
     assert store.get_unfinished_block2(unf4.partial_hash, None) == (unf1, 2, False)
 
-    ublock = store.get_unfinished_block_result(unf1.partial_hash)
-    assert ublock is not None and ublock.required_iters == uint64(0)
     ublock = store.get_unfinished_block_result2(unf1.partial_hash, unf1.foliage.foliage_transaction_block_hash)
     assert ublock is not None and ublock.required_iters == uint64(0)
-    # still, when not specifying a foliage hash, you just get the first ublock
-    ublock = store.get_unfinished_block_result2(unf1.partial_hash, None)
-    assert ublock is not None and ublock.required_iters == uint64(0)
 
     # negative test cases
-    assert store.get_unfinished_block_result(bytes32([1] * 32)) is None
-    assert store.get_unfinished_block_result2(bytes32([1] * 32), None) is None
+    assert store.get_unfinished_block_result2(bytes32([1] * 32), unf1.foliage.foliage_transaction_block_hash) is None
 
     blocks = custom_block_tools.get_consecutive_blocks(
         1,

tests/core/full_node/test_full_node.py

@@ -749,7 +749,9 @@ async def test_respond_unfinished(self, wallet_nodes, self_hostname):
         assert full_node_1.full_node.full_node_store.get_unfinished_block(unf.partial_hash) is None
         await full_node_1.full_node.add_unfinished_block(unf, None)
         assert full_node_1.full_node.full_node_store.get_unfinished_block(unf.partial_hash) is not None
-        result = full_node_1.full_node.full_node_store.get_unfinished_block_result(unf.partial_hash)
+        result = full_node_1.full_node.full_node_store.get_unfinished_block_result2(
+            unf.partial_hash, unf.foliage.foliage_transaction_block_hash
+        )
         assert result is not None
         assert result.npc_result is not None
         assert result.npc_result.conds is not None