refactor: move simulator utils functions (#772)

extras/custom_checks.sh

@@ -80,11 +80,21 @@ function check_deprecated_typing() {
 	return 0
 }
 
+function check_do_not_import_tests_in_hathor() {
+	if grep -R '\<.*import .*tests.*\>\|\<.*from .*tests.* import\>' "hathor"; then
+		echo 'do not import test definitions in the hathor module'
+		echo 'move them from tests to hathor instead'
+		return 1
+	fi
+	return 0
+}
+
 # List of functions to be executed
 checks=(
 	check_version_match
 	check_do_not_use_builtin_random_in_tests
 	check_deprecated_typing
+	check_do_not_import_tests_in_hathor
 )
 
 # Initialize a variable to track if any check fails

hathor/cli/events_simulator/scenario.py

@@ -44,14 +44,14 @@ def simulate_only_load(simulator: 'Simulator', _manager: 'HathorManager') -> Non
 
 
 def simulate_single_chain_one_block(simulator: 'Simulator', manager: 'HathorManager') -> None:
-    from tests.utils import add_new_blocks
+    from hathor.simulator.utils import add_new_blocks
     add_new_blocks(manager, 1)
     simulator.run(60)
 
 
 def simulate_single_chain_blocks_and_transactions(simulator: 'Simulator', manager: 'HathorManager') -> None:
     from hathor.conf.get_settings import get_settings
-    from tests.utils import add_new_blocks, gen_new_tx
+    from hathor.simulator.utils import add_new_blocks, gen_new_tx
 
     settings = get_settings()
     assert manager.wallet is not None
@@ -78,7 +78,7 @@ def simulate_single_chain_blocks_and_transactions(simulator: 'Simulator', manage
 
 def simulate_reorg(simulator: 'Simulator', manager: 'HathorManager') -> None:
     from hathor.simulator import FakeConnection
-    from tests.utils import add_new_blocks
+    from hathor.simulator.utils import add_new_blocks
 
     builder = simulator.get_default_builder()
     manager2 = simulator.create_peer(builder)

hathor/simulator/tx_generator.py

@@ -18,10 +18,10 @@
 from structlog import get_logger
 
 from hathor.conf.get_settings import get_settings
+from hathor.simulator.utils import NoCandidatesError, gen_new_double_spending, gen_new_tx
 from hathor.transaction.exceptions import RewardLocked
 from hathor.util import Random
 from hathor.wallet.exceptions import InsufficientFunds
-from tests.utils import NoCandidatesError, gen_new_double_spending, gen_new_tx
 
 if TYPE_CHECKING:
     from hathor.manager import HathorManager

hathor/simulator/utils.py

@@ -0,0 +1,183 @@
+#  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 typing import Optional, cast
+
+from hathor.crypto.util import decode_address
+from hathor.manager import HathorManager
+from hathor.transaction import Block, Transaction
+from hathor.types import Address, VertexId
+
+
+def gen_new_tx(manager: HathorManager, address: str, value: int, verify: bool = True) -> Transaction:
+    """
+    Generate and return a new transaction.
+
+    Args:
+        manager: the HathorManager to generate the transaction for
+        address: an address for the transaction's output
+        value: a value for the transaction's output
+        verify: whether to verify the generated transaction
+
+    Returns: the generated transaction.
+    """
+    from hathor.transaction import Transaction
+    from hathor.wallet.base_wallet import WalletOutputInfo
+
+    outputs = []
+    outputs.append(WalletOutputInfo(address=decode_address(address), value=int(value), timelock=None))
+
+    assert manager.wallet is not None
+    tx = manager.wallet.prepare_transaction_compute_inputs(Transaction, outputs, manager.tx_storage)
+    tx.storage = manager.tx_storage
+
+    max_ts_spent_tx = max(tx.get_spent_tx(txin).timestamp for txin in tx.inputs)
+    tx.timestamp = max(max_ts_spent_tx + 1, int(manager.reactor.seconds()))
+
+    tx.weight = 1
+    tx.parents = manager.get_new_tx_parents(tx.timestamp)
+    manager.cpu_mining_service.resolve(tx)
+    if verify:
+        manager.verification_service.verify(tx)
+    return tx
+
+
+def add_new_blocks(
+    manager: HathorManager,
+    num_blocks: int,
+    advance_clock: Optional[int] = None,
+    *,
+    parent_block_hash: Optional[VertexId] = None,
+    block_data: bytes = b'',
+    weight: Optional[float] = None,
+    address: Optional[Address] = None,
+    signal_bits: int | None = None,
+) -> list[Block]:
+    """ Create, resolve and propagate some blocks
+
+        :param manager: Manager object to handle the creation
+        :type manager: :py:class:`hathor.manager.HathorManager`
+
+        :param num_blocks: Quantity of blocks to be created
+        :type num_blocks: int
+
+        :return: Blocks created
+        :rtype: list[Block]
+    """
+    blocks = []
+    for _ in range(num_blocks):
+        blocks.append(
+            add_new_block(manager, advance_clock, parent_block_hash=parent_block_hash,
+                          data=block_data, weight=weight, address=address, signal_bits=signal_bits)
+        )
+        if parent_block_hash:
+            parent_block_hash = blocks[-1].hash
+    return blocks
+
+
+def add_new_block(
+    manager: HathorManager,
+    advance_clock: Optional[int] = None,
+    *,
+    parent_block_hash: Optional[VertexId] = None,
+    data: bytes = b'',
+    weight: Optional[float] = None,
+    address: Optional[Address] = None,
+    propagate: bool = True,
+    signal_bits: int | None = None,
+) -> Block:
+    """ Create, resolve and propagate a new block
+
+        :param manager: Manager object to handle the creation
+        :type manager: :py:class:`hathor.manager.HathorManager`
+
+        :return: Block created
+        :rtype: :py:class:`hathor.transaction.block.Block`
+    """
+    block = manager.generate_mining_block(parent_block_hash=parent_block_hash, data=data, address=address)
+    if weight is not None:
+        block.weight = weight
+    if signal_bits is not None:
+        block.signal_bits = signal_bits
+    manager.cpu_mining_service.resolve(block)
+    manager.verification_service.validate_full(block)
+    if propagate:
+        manager.propagate_tx(block, fails_silently=False)
+    if advance_clock:
+        assert hasattr(manager.reactor, 'advance')
+        manager.reactor.advance(advance_clock)
+    return block
+
+
+class NoCandidatesError(Exception):
+    pass
+
+
+def gen_new_double_spending(manager: HathorManager, *, use_same_parents: bool = False,
+                            tx: Optional[Transaction] = None, weight: float = 1) -> Transaction:
+    """
+    Generate and return a double spending transaction.
+
+    Args:
+        manager: the HathorManager to generate the transaction for
+        use_same_parents: whether to use the same parents as the original transaction
+        tx: the original transaction do double spend
+        weight: the new transaction's weight
+
+    Returns: the double spending transaction.
+    """
+    if tx is None:
+        tx_candidates = manager.get_new_tx_parents()
+        genesis = manager.tx_storage.get_all_genesis()
+        genesis_txs = [tx for tx in genesis if not tx.is_block]
+        # XXX: it isn't possible to double-spend a genesis transaction, thus we remove it from tx_candidates
+        for genesis_tx in genesis_txs:
+            if genesis_tx.hash in tx_candidates:
+                tx_candidates.remove(genesis_tx.hash)
+        if not tx_candidates:
+            raise NoCandidatesError()
+        # assert tx_candidates, 'Must not be empty, otherwise test was wrongly set up'
+        tx_hash = manager.rng.choice(tx_candidates)
+        tx = cast(Transaction, manager.tx_storage.get_transaction(tx_hash))
+
+    txin = manager.rng.choice(tx.inputs)
+
+    from hathor.transaction.scripts import P2PKH, parse_address_script
+    spent_tx = tx.get_spent_tx(txin)
+    spent_txout = spent_tx.outputs[txin.index]
+    p2pkh = parse_address_script(spent_txout.script)
+    assert isinstance(p2pkh, P2PKH)
+
+    from hathor.wallet.base_wallet import WalletInputInfo, WalletOutputInfo
+    value = spent_txout.value
+    wallet = manager.wallet
+    assert wallet is not None
+    private_key = wallet.get_private_key(p2pkh.address)
+    inputs = [WalletInputInfo(tx_id=txin.tx_id, index=txin.index, private_key=private_key)]
+
+    address = wallet.get_unused_address(mark_as_used=True)
+    outputs = [WalletOutputInfo(address=decode_address(address), value=int(value), timelock=None)]
+
+    tx2 = wallet.prepare_transaction(Transaction, inputs, outputs)
+    tx2.storage = manager.tx_storage
+    tx2.weight = weight
+    tx2.timestamp = max(tx.timestamp + 1, int(manager.reactor.seconds()))
+
+    if use_same_parents:
+        tx2.parents = list(tx.parents)
+    else:
+        tx2.parents = manager.get_new_tx_parents(tx2.timestamp)
+
+    manager.cpu_mining_service.resolve(tx2)
+    return tx2

tests/cli/test_multisig_signature.py

@@ -9,9 +9,10 @@
 from structlog.testing import capture_logs
 
 from hathor.cli.multisig_signature import create_parser, execute
+from hathor.simulator.utils import add_new_blocks
 from hathor.wallet import Wallet
 from tests import unittest
-from tests.utils import add_blocks_unlock_reward, add_new_blocks, add_new_transactions
+from tests.utils import add_blocks_unlock_reward, add_new_transactions
 
 
 class BaseSignatureTest(unittest.TestCase):

tests/cli/test_multisig_spend.py

@@ -6,12 +6,13 @@
 from hathor.cli.multisig_spend import create_parser, execute
 from hathor.conf import HathorSettings
 from hathor.crypto.util import decode_address
+from hathor.simulator.utils import add_new_blocks
 from hathor.transaction import Transaction, TxInput, TxOutput
 from hathor.transaction.scripts import create_output_script
 from hathor.wallet.base_wallet import WalletBalance, WalletOutputInfo
 from hathor.wallet.util import generate_multisig_address, generate_multisig_redeem_script, generate_signature
 from tests import unittest
-from tests.utils import add_blocks_unlock_reward, add_new_blocks
+from tests.utils import add_blocks_unlock_reward
 
 settings = HathorSettings()
 

tests/cli/test_twin_tx.py

@@ -6,12 +6,12 @@
 
 from hathor.cli.twin_tx import create_parser, execute
 from hathor.conf import HathorSettings
+from hathor.simulator.utils import add_new_blocks
 from hathor.transaction import Transaction, TransactionMetadata
 from hathor.util import json_loadb
 from tests import unittest
 from tests.utils import (
     add_blocks_unlock_reward,
-    add_new_blocks,
     add_new_transactions,
     execute_mining,
     execute_tx_gen,

tests/consensus/test_consensus.py

@@ -1,16 +1,10 @@
 from unittest.mock import MagicMock
 
 from hathor.conf import HathorSettings
+from hathor.simulator.utils import add_new_block, add_new_blocks, gen_new_tx
 from hathor.transaction.storage import TransactionMemoryStorage
 from tests import unittest
-from tests.utils import (
-    add_blocks_unlock_reward,
-    add_new_block,
-    add_new_blocks,
-    add_new_double_spending,
-    add_new_transactions,
-    gen_new_tx,
-)
+from tests.utils import add_blocks_unlock_reward, add_new_double_spending, add_new_transactions
 
 settings = HathorSettings()
 

tests/consensus/test_consensus2.py

@@ -1,7 +1,8 @@
 from hathor.graphviz import GraphvizVisualizer
+from hathor.simulator.utils import gen_new_tx
 from tests import unittest
 from tests.simulation.base import SimulatorTestCase
-from tests.utils import add_custom_tx, gen_new_tx
+from tests.utils import add_custom_tx
 
 
 class BaseConsensusSimulatorTestCase(SimulatorTestCase):

tests/consensus/test_consensus3.py

@@ -1,7 +1,8 @@
 import pytest
 
+from hathor.simulator.utils import add_new_block, add_new_blocks
 from tests import unittest
-from tests.utils import add_blocks_unlock_reward, add_new_block, add_new_blocks
+from tests.utils import add_blocks_unlock_reward
 
 
 class DoubleSpendingTestCase(unittest.TestCase):

tests/consensus/test_soft_voided.py

@@ -2,9 +2,10 @@
 from hathor.graphviz import GraphvizVisualizer
 from hathor.simulator import FakeConnection, Simulator
 from hathor.simulator.trigger import StopAfterNTransactions
+from hathor.simulator.utils import gen_new_tx
 from tests import unittest
 from tests.simulation.base import SimulatorTestCase
-from tests.utils import add_custom_tx, gen_new_tx
+from tests.utils import add_custom_tx
 
 settings = HathorSettings()
 

tests/consensus/test_soft_voided2.py

@@ -1,9 +1,10 @@
 from hathor.conf import HathorSettings
 from hathor.graphviz import GraphvizVisualizer
 from hathor.simulator import Simulator
+from hathor.simulator.utils import gen_new_tx
 from tests import unittest
 from tests.simulation.base import SimulatorTestCase
-from tests.utils import BURN_ADDRESS, add_custom_tx, gen_new_tx
+from tests.utils import BURN_ADDRESS, add_custom_tx
 
 settings = HathorSettings()
 

tests/consensus/test_soft_voided3.py

@@ -2,9 +2,10 @@
 from hathor.graphviz import GraphvizVisualizer
 from hathor.simulator import FakeConnection, Simulator
 from hathor.simulator.trigger import StopAfterNTransactions
+from hathor.simulator.utils import gen_new_tx
 from tests import unittest
 from tests.simulation.base import SimulatorTestCase
-from tests.utils import add_custom_tx, gen_custom_tx, gen_new_tx
+from tests.utils import add_custom_tx, gen_custom_tx
 
 settings = HathorSettings()
 

tests/consensus/test_soft_voided4.py

@@ -2,9 +2,10 @@
 from hathor.graphviz import GraphvizVisualizer
 from hathor.simulator import FakeConnection, Simulator
 from hathor.simulator.trigger import StopAfterNTransactions
+from hathor.simulator.utils import gen_new_double_spending
 from tests import unittest
 from tests.simulation.base import SimulatorTestCase
-from tests.utils import add_custom_tx, gen_new_double_spending
+from tests.utils import add_custom_tx
 
 settings = HathorSettings()
 

tests/event/test_event_reorg.py

@@ -1,8 +1,9 @@
 from hathor.conf import HathorSettings
 from hathor.event.model.event_type import EventType
 from hathor.event.storage import EventMemoryStorage
+from hathor.simulator.utils import add_new_blocks
 from tests import unittest
-from tests.utils import BURN_ADDRESS, add_new_blocks, get_genesis_key
+from tests.utils import BURN_ADDRESS, get_genesis_key
 
 settings = HathorSettings()