chore: improve miner simulator flakiness

hathor/simulator/miner/geometric_miner.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import math
 from typing import TYPE_CHECKING, Optional
 
 from hathor.conf import HathorSettings
@@ -50,6 +51,7 @@ def __init__(
         self._signal_bits = signal_bits or []
         self._block: Optional[Block] = None
         self._blocks_found: int = 0
+        self._blocks_before_pause: float = math.inf
 
     def _on_new_tx(self, key: HathorEvents, args: 'EventArguments') -> None:
         """ Called when a new tx or block is received. It updates the current mining to the
@@ -81,12 +83,17 @@ def _generate_mining_block(self) -> 'Block':
         return block
 
     def _schedule_next_block(self):
+        if self._blocks_before_pause <= 0:
+            self._delayed_call = None
+            return
+
         if self._block:
             self._block.nonce = self._rng.getrandbits(32)
             self._block.update_hash()
             self.log.debug('randomized step: found new block', hash=self._block.hash_hex, nonce=self._block.nonce)
             self._manager.propagate_tx(self._block, fails_silently=False)
             self._blocks_found += 1
+            self._blocks_before_pause -= 1
             self._block = None
 
         if self._manager.can_start_mining():
@@ -110,3 +117,16 @@ def _schedule_next_block(self):
 
     def get_blocks_found(self) -> int:
         return self._blocks_found
+
+    def pause_after_exactly(self, *, n_blocks: int) -> None:
+        """
+        Configure the miner to pause mining blocks after exactly `n_blocks` are propagated. If called more than once,
+        will unpause the miner and pause again according to the new argument.
+
+        Use this instead of the `StopAfterNMinedBlocks` trigger if you need "exactly N blocks" behavior, instead of
+        "at least N blocks".
+        """
+        self._blocks_before_pause = n_blocks
+
+        if not self._delayed_call:
+            self._delayed_call = self._clock.callLater(0, self._schedule_next_block)

hathor/simulator/trigger.py

@@ -31,7 +31,12 @@ def should_stop(self) -> bool:
 
 
 class StopAfterNMinedBlocks(Trigger):
-    """Stop the simulation after `miner` finds N blocks. Note that these blocks might be orphan."""
+    """
+    Stop the simulation after `miner` finds at least N blocks. Note that these blocks might be orphan.
+
+    Use `miner.pause_after_exactly()` instead of this trigger if you need "exactly N blocks" behavior, instead of
+    "at least N blocks".
+    """
     def __init__(self, miner: 'AbstractMiner', *, quantity: int) -> None:
         self.miner = miner
         self.quantity = quantity

tests/feature_activation/test_mining_simulation.py

@@ -26,7 +26,6 @@
 from hathor.feature_activation.settings import Settings as FeatureSettings
 from hathor.mining.ws import MiningWebsocketFactory, MiningWebsocketProtocol
 from hathor.p2p.resources import MiningResource
-from hathor.simulator.trigger import StopAfterNMinedBlocks
 from hathor.transaction.resources import GetBlockTemplateResource
 from hathor.transaction.util import unpack, unpack_len
 from hathor.util import json_loadb
@@ -88,47 +87,56 @@ def test_signal_bits_in_mining(self) -> None:
         # At the beginning, all features are outside their signaling period, so none are signaled.
         expected_signal_bits = 0b0000
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits]
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1))
+        miner.pause_after_exactly(n_blocks=1)
+        self.simulator.run(3600)
         assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits
         assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits]
 
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=6))
+        miner.pause_after_exactly(n_blocks=6)
+        self.simulator.run(3600)
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 6
 
         # At height=8, NOP_FEATURE_1 is signaling, so it's enabled by the default support.
         expected_signal_bits = 0b0001
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1))
+        miner.pause_after_exactly(n_blocks=1)
+        self.simulator.run(3600)
         assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits
         assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits]
 
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=3))
+        miner.pause_after_exactly(n_blocks=3)
+        self.simulator.run(3600)
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 3
 
         # At height=12, NOP_FEATURE_2 is signaling, enabled by the user. NOP_FEATURE_1 also continues signaling.
         expected_signal_bits = 0b0101
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1))
+        miner.pause_after_exactly(n_blocks=1)
+        self.simulator.run(3600)
         assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits
         assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits]
 
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=7))
+        miner.pause_after_exactly(n_blocks=7)
+        self.simulator.run(3600)
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 7
 
         # At height=20, NOP_FEATURE_1 stops signaling, and NOP_FEATURE_2 continues.
         expected_signal_bits = 0b0100
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1))
+        miner.pause_after_exactly(n_blocks=1)
+        self.simulator.run(3600)
         assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits
         assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits]
 
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=3))
+        miner.pause_after_exactly(n_blocks=3)
+        self.simulator.run(3600)
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits] * 3
 
         # At height=24, all features have left their signaling period and therefore none are signaled.
         expected_signal_bits = 0b0000
-        self.simulator.run(3600, trigger=StopAfterNMinedBlocks(miner, quantity=1))
+        miner.pause_after_exactly(n_blocks=1)
+        self.simulator.run(3600)
         assert self._get_signal_bits_from_get_block_template(get_block_template_client) == expected_signal_bits
         assert self._get_signal_bits_from_mining(mining_client) == expected_signal_bits
         assert self._get_ws_signal_bits(ws_transport) == [expected_signal_bits]