fix block_tools feature when specifying a list of block references. A…

…lso add feature keep_going_until_tx_block. (#11185)
Chia-Network · Apr 15, 2022 · 36a610f · 36a610f
1 parent 82c01ff
commit 36a610f
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/tests/block_tools.py b/tests/block_tools.py
@@ -419,6 +419,7 @@ def get_consecutive_blocks(
         self,
         num_blocks: int,
         block_list_input: List[FullBlock] = None,
+        *,
         farmer_reward_puzzle_hash: Optional[bytes32] = None,
         pool_reward_puzzle_hash: Optional[bytes32] = None,
         transaction_data: Optional[SpendBundle] = None,
@@ -427,6 +428,7 @@ def get_consecutive_blocks(
         force_overflow: bool = False,
         skip_slots: int = 0,  # Force at least this number of empty slots before the first SB
         guarantee_transaction_block: bool = False,  # Force that this block must be a tx block
+        keep_going_until_tx_block: bool = False,  # keep making new blocks until we find a tx block
         normalized_to_identity_cc_eos: bool = False,
         normalized_to_identity_icc_eos: bool = False,
         normalized_to_identity_cc_sp: bool = False,
@@ -566,7 +568,8 @@ def get_consecutive_blocks(
                         removals = None
                         if transaction_data_included:
                             transaction_data = None
-                        if transaction_data is not None and not transaction_data_included:
+                            previous_generator = None
+                        if transaction_data is not None:
                             additions = transaction_data.additions()
                             removals = transaction_data.removals()
                         assert start_timestamp is not None
@@ -582,9 +585,10 @@ def get_consecutive_blocks(
                             else:
                                 pool_target = PoolTarget(self.pool_ph, uint32(0))
 
+                        block_generator: Optional[BlockGenerator]
                         if transaction_data is not None:
                             if type(previous_generator) is CompressorArg:
-                                block_generator: Optional[BlockGenerator] = best_solution_generator_from_template(
+                                block_generator = best_solution_generator_from_template(
                                     previous_generator, transaction_data
                                 )
                             else:
@@ -629,6 +633,8 @@ def get_consecutive_blocks(
                         )
                         if block_record.is_transaction_block:
                             transaction_data_included = True
+                            previous_generator = None
+                            keep_going_until_tx_block = False
                         else:
                             if guarantee_transaction_block:
                                 continue
@@ -650,7 +656,7 @@ def get_consecutive_blocks(
                         latest_block = blocks[full_block.header_hash]
                         finished_sub_slots_at_ip = []
                         num_blocks -= 1
-                        if num_blocks == 0:
+                        if num_blocks <= 0 and not keep_going_until_tx_block:
                             return block_list
 
             # Finish the end of sub-slot and try again next sub-slot
@@ -789,7 +795,7 @@ def get_consecutive_blocks(
             removals = None
             if transaction_data_included:
                 transaction_data = None
-            if transaction_data is not None and not transaction_data_included:
+            if transaction_data is not None:
                 additions = transaction_data.additions()
                 removals = transaction_data.removals()
             sub_slots_finished += 1
@@ -856,6 +862,8 @@ def get_consecutive_blocks(
                                 )
                             else:
                                 block_generator = simple_solution_generator(transaction_data)
+                                if type(previous_generator) is list:
+                                    block_generator = BlockGenerator(block_generator.program, [], previous_generator)
                             aggregate_signature = transaction_data.aggregated_signature
                         else:
                             block_generator = None
@@ -895,6 +903,8 @@ def get_consecutive_blocks(
 
                         if block_record.is_transaction_block:
                             transaction_data_included = True
+                            previous_generator = None
+                            keep_going_until_tx_block = False
                         elif guarantee_transaction_block:
                             continue
                         if pending_ses:
@@ -911,7 +921,7 @@ def get_consecutive_blocks(
                         blocks_added_this_sub_slot += 1
                         log.info(f"Created block {block_record.height } ov=True, iters " f"{block_record.total_iters}")
                         num_blocks -= 1
-                        if num_blocks == 0:
+                        if num_blocks <= 0 and not keep_going_until_tx_block:
                             return block_list
 
                         blocks[full_block.header_hash] = block_record

diff --git a/tests/blockchain/test_blockchain.py b/tests/blockchain/test_blockchain.py
@@ -3280,7 +3280,13 @@ async def test_reorg_flip_flop(empty_blockchain, bt):
     )
 
     spend_bundle = wallet_a.generate_signed_transaction(1000, receiver_puzzlehash, all_coins.pop())
-    chain_a = bt.get_consecutive_blocks(5, chain_a, previous_generator=[uint32(10)], transaction_data=spend_bundle)
+    chain_a = bt.get_consecutive_blocks(
+        5,
+        chain_a,
+        previous_generator=[uint32(10)],
+        transaction_data=spend_bundle,
+        guarantee_transaction_block=True,
+    )
 
     spend_bundle = wallet_a.generate_signed_transaction(1000, receiver_puzzlehash, all_coins.pop())
     chain_a = bt.get_consecutive_blocks(