feat(simulator): Add a trigger to stop the simulator when a condition…

… is satisfied

hathor/simulator/miner/__init__.py

@@ -0,0 +1,21 @@
+#  Copyright 2023 Hathor Labs
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#  http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from hathor.simulator.miner.abstract_miner import AbstractMiner
+from hathor.simulator.miner.geometric_miner import GeometricMiner
+
+__all__ = [
+    'AbstractMiner',
+    'GeometricMiner',
+]

hathor/simulator/miner/abstract_miner.py

@@ -63,3 +63,7 @@ def _on_new_tx(self, key: HathorEvents, args: EventArguments) -> None:
     def _schedule_next_block(self):
         """Schedule the propagation of the next block, and propagate a block if it has been found."""
         raise NotImplementedError
+
+    @abstractmethod
+    def get_blocks_found(self) -> int:
+        raise NotImplementedError

hathor/simulator/miner/geometric_miner.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from hathor.conf import HathorSettings
 from hathor.manager import HathorEvents
@@ -22,6 +22,7 @@
 if TYPE_CHECKING:
     from hathor.manager import HathorManager
     from hathor.pubsub import EventArguments
+    from hathor.transaction import Block
 
 settings = HathorSettings()
 
@@ -37,7 +38,8 @@ def __init__(self, manager: 'HathorManager', rng: Random, *, hashpower: float):
         super().__init__(manager, rng)
 
         self._hashpower = hashpower
-        self._block = None
+        self._block: Optional[Block] = None
+        self._blocks_found: int = 0
 
     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
@@ -49,6 +51,7 @@ def _on_new_tx(self, key: HathorEvents, args: 'EventArguments') -> None:
         if not self._block:
             return
 
+        assert tx.storage is not None
         tips = tx.storage.get_best_block_tips()
         if self._block.parents[0] not in tips:
             # Head changed
@@ -61,6 +64,7 @@ def _schedule_next_block(self):
             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._block = None
 
         if self._manager.can_start_mining():
@@ -81,3 +85,6 @@ def _schedule_next_block(self):
         if self._delayed_call and self._delayed_call.active():
             self._delayed_call.cancel()
         self._delayed_call = self._clock.callLater(dt, self._schedule_next_block)
+
+    def get_blocks_found(self) -> int:
+        return self._blocks_found

hathor/simulator/simulator.py

@@ -21,6 +21,7 @@
 from structlog import get_logger
 
 from hathor.builder import Builder
+from hathor.conf import HathorSettings
 from hathor.daa import TestMode, _set_test_mode
 from hathor.manager import HathorManager
 from hathor.p2p.peer_id import PeerId
@@ -34,6 +35,7 @@
 
 if TYPE_CHECKING:
     from hathor.simulator.fake_connection import FakeConnection
+    from hathor.simulator.trigger import Trigger
 
 
 logger = get_logger()
@@ -104,6 +106,7 @@ def __init__(self, seed: Optional[int] = None):
             seed = secrets.randbits(64)
         self.seed = seed
         self.rng = Random(self.seed)
+        self.settings = HathorSettings()
         self._network = 'testnet'
         self._clock = HeapClock()
         self._peers: OrderedDict[str, HathorManager] = OrderedDict()
@@ -186,6 +189,9 @@ def add_peer(self, name: str, peer: HathorManager) -> None:
             raise ValueError('Duplicate peer name')
         self._peers[name] = peer
 
+    def get_reactor(self) -> HeapClock:
+        return self._clock
+
     def get_peer(self, name: str) -> HathorManager:
         return self._peers[name]
 
@@ -242,7 +248,18 @@ def run_until_complete(self,
     def run(self,
             interval: float,
             step: float = DEFAULT_STEP_INTERVAL,
-            status_interval: float = DEFAULT_STATUS_INTERVAL) -> None:
+            status_interval: float = DEFAULT_STATUS_INTERVAL,
+            *,
+            trigger: Optional['Trigger'] = None) -> bool:
+        """Return True if it successfully ends the execution.
+
+        If no trigger is provided, it always returns True.
+        If a trigger is provided, it returns True if the trigger stops the execution. Otherwise, it returns False.
+        """
         assert self._started
         for _ in self._run(interval, step, status_interval):
-            pass
+            if trigger is not None and trigger.should_stop():
+                return True
+        if trigger is not None:
+            return False
+        return True

hathor/simulator/trigger.py

@@ -0,0 +1,56 @@
+# Copyright 2023 Hathor Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from hathor.simulator.miner import AbstractMiner
+    from hathor.wallet import BaseWallet
+
+
+class Trigger(ABC):
+    """Abstract class to stop simulation when a certain condition is satisfied."""
+    @abstractmethod
+    def should_stop(self) -> bool:
+        """This method must return True when the stop condition is satisfied."""
+        raise NotImplementedError
+
+
+class StopAfterNMinedBlocks(Trigger):
+    """Stop the simulation after `miner` finds N blocks. Note that these blocks might be orphan."""
+    def __init__(self, miner: 'AbstractMiner', *, quantity: int) -> None:
+        self.miner = miner
+        self.quantity = quantity
+        self.reset()
+
+    def reset(self) -> None:
+        """Reset the counter, so this trigger can be reused."""
+        self.initial_blocks_found = self.miner.get_blocks_found()
+
+    def should_stop(self) -> bool:
+        diff = self.miner.get_blocks_found() - self.initial_blocks_found
+        return diff >= self.quantity
+
+
+class StopAfterMinimumBalance(Trigger):
+    """Stop the simulation after `wallet` reaches a minimum unlocked balance."""
+    def __init__(self, wallet: 'BaseWallet', token_uid: bytes, minimum_balance: int) -> None:
+        self.wallet = wallet
+        self.token_uid = token_uid
+        self.minimum_balance = minimum_balance
+
+    def should_stop(self) -> bool:
+        balance = self.wallet.balance[self.token_uid].available
+        return balance >= self.minimum_balance

tests/simulation/test_trigger.py

@@ -0,0 +1,58 @@
+from hathor.p2p.peer_id import PeerId
+from hathor.simulator import Simulator
+from hathor.simulator.trigger import StopAfterMinimumBalance, StopAfterNMinedBlocks
+from tests import unittest
+
+
+class TriggerTestCase(unittest.TestCase):
+    def setUp(self):
+        super().setUp()
+
+        self.simulator = Simulator()
+        self.simulator.start()
+
+        peer_id = PeerId()
+        self.manager1 = self.simulator.create_peer(peer_id=peer_id)
+        self.manager1.allow_mining_without_peers()
+
+        print('-' * 30)
+        print('Simulation seed config:', self.simulator.seed)
+        print('-' * 30)
+
+    def test_stop_after_n_mined_blocks(self):
+        miner1 = self.simulator.create_miner(self.manager1, hashpower=1e6)
+        miner1.start()
+
+        reactor = self.simulator.get_reactor()
+
+        t0 = reactor.seconds()
+        trigger = StopAfterNMinedBlocks(miner1, quantity=3)
+        self.assertEqual(miner1.get_blocks_found(), 0)
+        self.assertTrue(self.simulator.run(3600, trigger=trigger))
+        self.assertEqual(miner1.get_blocks_found(), 3)
+        self.assertLess(reactor.seconds(), t0 + 3600)
+
+        trigger.reset()
+        self.assertTrue(self.simulator.run(3600, trigger=trigger))
+        self.assertEqual(miner1.get_blocks_found(), 6)
+
+        t0 = reactor.seconds()
+        trigger = StopAfterNMinedBlocks(miner1, quantity=10)
+        self.assertTrue(self.simulator.run(3600, trigger=trigger))
+        self.assertEqual(miner1.get_blocks_found(), 16)
+        self.assertLess(reactor.seconds(), t0 + 3600)
+
+    def test_stop_after_minimum_balance(self):
+        miner1 = self.simulator.create_miner(self.manager1, hashpower=1e6)
+        miner1.start()
+
+        wallet = self.manager1.wallet
+        settings = self.simulator.settings
+
+        minimum_balance = 1000_00   # 16 blocks
+        token_uid = settings.HATHOR_TOKEN_UID
+
+        trigger = StopAfterMinimumBalance(wallet, token_uid, minimum_balance)
+        self.assertLess(wallet.balance[token_uid].available, minimum_balance)
+        self.assertTrue(self.simulator.run(3600, trigger=trigger))
+        self.assertGreaterEqual(wallet.balance[token_uid].available, minimum_balance)