Skip to content

Commit cacba12

Browse files
committed
refactor: move simulator utils functions
1 parent 769d255 commit cacba12

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+297
-214
lines changed

extras/custom_checks.sh

+10
Original file line numberDiff line numberDiff line change
@@ -80,11 +80,21 @@ function check_deprecated_typing() {
8080
return 0
8181
}
8282

83+
function check_do_not_import_tests_in_hathor() {
84+
if grep -R '\<.*import .*tests.*\>\|\<.*from .*tests.* import\>' "hathor"; then
85+
echo 'do not import test definitions in the hathor module'
86+
echo 'move them from tests to hathor instead'
87+
return 1
88+
fi
89+
return 0
90+
}
91+
8392
# List of functions to be executed
8493
checks=(
8594
check_version_match
8695
check_do_not_use_builtin_random_in_tests
8796
check_deprecated_typing
97+
check_do_not_import_tests_in_hathor
8898
)
8999

90100
# Initialize a variable to track if any check fails

hathor/cli/events_simulator/scenario.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -44,14 +44,14 @@ def simulate_only_load(simulator: 'Simulator', _manager: 'HathorManager') -> Non
4444

4545

4646
def simulate_single_chain_one_block(simulator: 'Simulator', manager: 'HathorManager') -> None:
47-
from tests.utils import add_new_blocks
47+
from hathor.simulator.utils import add_new_blocks
4848
add_new_blocks(manager, 1)
4949
simulator.run(60)
5050

5151

5252
def simulate_single_chain_blocks_and_transactions(simulator: 'Simulator', manager: 'HathorManager') -> None:
5353
from hathor.conf.get_settings import get_settings
54-
from tests.utils import add_new_blocks, gen_new_tx
54+
from hathor.simulator.utils import add_new_blocks, gen_new_tx
5555

5656
settings = get_settings()
5757
assert manager.wallet is not None
@@ -78,7 +78,7 @@ def simulate_single_chain_blocks_and_transactions(simulator: 'Simulator', manage
7878

7979
def simulate_reorg(simulator: 'Simulator', manager: 'HathorManager') -> None:
8080
from hathor.simulator import FakeConnection
81-
from tests.utils import add_new_blocks
81+
from hathor.simulator.utils import add_new_blocks
8282

8383
builder = simulator.get_default_builder()
8484
manager2 = simulator.create_peer(builder)

hathor/simulator/tx_generator.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,10 @@
1818
from structlog import get_logger
1919

2020
from hathor.conf.get_settings import get_settings
21+
from hathor.simulator.utils import NoCandidatesError, gen_new_double_spending, gen_new_tx
2122
from hathor.transaction.exceptions import RewardLocked
2223
from hathor.util import Random
2324
from hathor.wallet.exceptions import InsufficientFunds
24-
from tests.utils import NoCandidatesError, gen_new_double_spending, gen_new_tx
2525

2626
if TYPE_CHECKING:
2727
from hathor.manager import HathorManager

hathor/simulator/utils.py

+183
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
# Copyright 2023 Hathor Labs
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from typing import Optional, cast
16+
17+
from hathor.crypto.util import decode_address
18+
from hathor.manager import HathorManager
19+
from hathor.transaction import Block, Transaction
20+
from hathor.types import Address, VertexId
21+
22+
23+
def gen_new_tx(manager: HathorManager, address: str, value: int, verify: bool = True) -> Transaction:
24+
"""
25+
Generate and return a new transaction.
26+
27+
Args:
28+
manager: the HathorManager to generate the transaction for
29+
address: an address for the transaction's output
30+
value: a value for the transaction's output
31+
verify: whether to verify the generated transaction
32+
33+
Returns: the generated transaction.
34+
"""
35+
from hathor.transaction import Transaction
36+
from hathor.wallet.base_wallet import WalletOutputInfo
37+
38+
outputs = []
39+
outputs.append(WalletOutputInfo(address=decode_address(address), value=int(value), timelock=None))
40+
41+
assert manager.wallet is not None
42+
tx = manager.wallet.prepare_transaction_compute_inputs(Transaction, outputs, manager.tx_storage)
43+
tx.storage = manager.tx_storage
44+
45+
max_ts_spent_tx = max(tx.get_spent_tx(txin).timestamp for txin in tx.inputs)
46+
tx.timestamp = max(max_ts_spent_tx + 1, int(manager.reactor.seconds()))
47+
48+
tx.weight = 1
49+
tx.parents = manager.get_new_tx_parents(tx.timestamp)
50+
manager.cpu_mining_service.resolve(tx)
51+
if verify:
52+
manager.verification_service.verify(tx)
53+
return tx
54+
55+
56+
def add_new_blocks(
57+
manager: HathorManager,
58+
num_blocks: int,
59+
advance_clock: Optional[int] = None,
60+
*,
61+
parent_block_hash: Optional[VertexId] = None,
62+
block_data: bytes = b'',
63+
weight: Optional[float] = None,
64+
address: Optional[Address] = None,
65+
signal_bits: int | None = None,
66+
) -> list[Block]:
67+
""" Create, resolve and propagate some blocks
68+
69+
:param manager: Manager object to handle the creation
70+
:type manager: :py:class:`hathor.manager.HathorManager`
71+
72+
:param num_blocks: Quantity of blocks to be created
73+
:type num_blocks: int
74+
75+
:return: Blocks created
76+
:rtype: list[Block]
77+
"""
78+
blocks = []
79+
for _ in range(num_blocks):
80+
blocks.append(
81+
add_new_block(manager, advance_clock, parent_block_hash=parent_block_hash,
82+
data=block_data, weight=weight, address=address, signal_bits=signal_bits)
83+
)
84+
if parent_block_hash:
85+
parent_block_hash = blocks[-1].hash
86+
return blocks
87+
88+
89+
def add_new_block(
90+
manager: HathorManager,
91+
advance_clock: Optional[int] = None,
92+
*,
93+
parent_block_hash: Optional[VertexId] = None,
94+
data: bytes = b'',
95+
weight: Optional[float] = None,
96+
address: Optional[Address] = None,
97+
propagate: bool = True,
98+
signal_bits: int | None = None,
99+
) -> Block:
100+
""" Create, resolve and propagate a new block
101+
102+
:param manager: Manager object to handle the creation
103+
:type manager: :py:class:`hathor.manager.HathorManager`
104+
105+
:return: Block created
106+
:rtype: :py:class:`hathor.transaction.block.Block`
107+
"""
108+
block = manager.generate_mining_block(parent_block_hash=parent_block_hash, data=data, address=address)
109+
if weight is not None:
110+
block.weight = weight
111+
if signal_bits is not None:
112+
block.signal_bits = signal_bits
113+
manager.cpu_mining_service.resolve(block)
114+
manager.verification_service.validate_full(block)
115+
if propagate:
116+
manager.propagate_tx(block, fails_silently=False)
117+
if advance_clock:
118+
assert hasattr(manager.reactor, 'advance')
119+
manager.reactor.advance(advance_clock)
120+
return block
121+
122+
123+
class NoCandidatesError(Exception):
124+
pass
125+
126+
127+
def gen_new_double_spending(manager: HathorManager, *, use_same_parents: bool = False,
128+
tx: Optional[Transaction] = None, weight: float = 1) -> Transaction:
129+
"""
130+
Generate and return a double spending transaction.
131+
132+
Args:
133+
manager: the HathorManager to generate the transaction for
134+
use_same_parents: whether to use the same parents as the original transaction
135+
tx: the original transaction do double spend
136+
weight: the new transaction's weight
137+
138+
Returns: the double spending transaction.
139+
"""
140+
if tx is None:
141+
tx_candidates = manager.get_new_tx_parents()
142+
genesis = manager.tx_storage.get_all_genesis()
143+
genesis_txs = [tx for tx in genesis if not tx.is_block]
144+
# XXX: it isn't possible to double-spend a genesis transaction, thus we remove it from tx_candidates
145+
for genesis_tx in genesis_txs:
146+
if genesis_tx.hash in tx_candidates:
147+
tx_candidates.remove(genesis_tx.hash)
148+
if not tx_candidates:
149+
raise NoCandidatesError()
150+
# assert tx_candidates, 'Must not be empty, otherwise test was wrongly set up'
151+
tx_hash = manager.rng.choice(tx_candidates)
152+
tx = cast(Transaction, manager.tx_storage.get_transaction(tx_hash))
153+
154+
txin = manager.rng.choice(tx.inputs)
155+
156+
from hathor.transaction.scripts import P2PKH, parse_address_script
157+
spent_tx = tx.get_spent_tx(txin)
158+
spent_txout = spent_tx.outputs[txin.index]
159+
p2pkh = parse_address_script(spent_txout.script)
160+
assert isinstance(p2pkh, P2PKH)
161+
162+
from hathor.wallet.base_wallet import WalletInputInfo, WalletOutputInfo
163+
value = spent_txout.value
164+
wallet = manager.wallet
165+
assert wallet is not None
166+
private_key = wallet.get_private_key(p2pkh.address)
167+
inputs = [WalletInputInfo(tx_id=txin.tx_id, index=txin.index, private_key=private_key)]
168+
169+
address = wallet.get_unused_address(mark_as_used=True)
170+
outputs = [WalletOutputInfo(address=decode_address(address), value=int(value), timelock=None)]
171+
172+
tx2 = wallet.prepare_transaction(Transaction, inputs, outputs)
173+
tx2.storage = manager.tx_storage
174+
tx2.weight = weight
175+
tx2.timestamp = max(tx.timestamp + 1, int(manager.reactor.seconds()))
176+
177+
if use_same_parents:
178+
tx2.parents = list(tx.parents)
179+
else:
180+
tx2.parents = manager.get_new_tx_parents(tx2.timestamp)
181+
182+
manager.cpu_mining_service.resolve(tx2)
183+
return tx2

tests/cli/test_multisig_signature.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,10 @@
99
from structlog.testing import capture_logs
1010

1111
from hathor.cli.multisig_signature import create_parser, execute
12+
from hathor.simulator.utils import add_new_blocks
1213
from hathor.wallet import Wallet
1314
from tests import unittest
14-
from tests.utils import add_blocks_unlock_reward, add_new_blocks, add_new_transactions
15+
from tests.utils import add_blocks_unlock_reward, add_new_transactions
1516

1617

1718
class BaseSignatureTest(unittest.TestCase):

tests/cli/test_multisig_spend.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,13 @@
66
from hathor.cli.multisig_spend import create_parser, execute
77
from hathor.conf import HathorSettings
88
from hathor.crypto.util import decode_address
9+
from hathor.simulator.utils import add_new_blocks
910
from hathor.transaction import Transaction, TxInput, TxOutput
1011
from hathor.transaction.scripts import create_output_script
1112
from hathor.wallet.base_wallet import WalletBalance, WalletOutputInfo
1213
from hathor.wallet.util import generate_multisig_address, generate_multisig_redeem_script, generate_signature
1314
from tests import unittest
14-
from tests.utils import add_blocks_unlock_reward, add_new_blocks
15+
from tests.utils import add_blocks_unlock_reward
1516

1617
settings = HathorSettings()
1718

tests/cli/test_twin_tx.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@
66

77
from hathor.cli.twin_tx import create_parser, execute
88
from hathor.conf import HathorSettings
9+
from hathor.simulator.utils import add_new_blocks
910
from hathor.transaction import Transaction, TransactionMetadata
1011
from hathor.util import json_loadb
1112
from tests import unittest
1213
from tests.utils import (
1314
add_blocks_unlock_reward,
14-
add_new_blocks,
1515
add_new_transactions,
1616
execute_mining,
1717
execute_tx_gen,

tests/consensus/test_consensus.py

+2-8
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,10 @@
11
from unittest.mock import MagicMock
22

33
from hathor.conf import HathorSettings
4+
from hathor.simulator.utils import add_new_block, add_new_blocks, gen_new_tx
45
from hathor.transaction.storage import TransactionMemoryStorage
56
from tests import unittest
6-
from tests.utils import (
7-
add_blocks_unlock_reward,
8-
add_new_block,
9-
add_new_blocks,
10-
add_new_double_spending,
11-
add_new_transactions,
12-
gen_new_tx,
13-
)
7+
from tests.utils import add_blocks_unlock_reward, add_new_double_spending, add_new_transactions
148

159
settings = HathorSettings()
1610

tests/consensus/test_consensus2.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
from hathor.graphviz import GraphvizVisualizer
2+
from hathor.simulator.utils import gen_new_tx
23
from tests import unittest
34
from tests.simulation.base import SimulatorTestCase
4-
from tests.utils import add_custom_tx, gen_new_tx
5+
from tests.utils import add_custom_tx
56

67

78
class BaseConsensusSimulatorTestCase(SimulatorTestCase):

tests/consensus/test_consensus3.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
import pytest
22

3+
from hathor.simulator.utils import add_new_block, add_new_blocks
34
from tests import unittest
4-
from tests.utils import add_blocks_unlock_reward, add_new_block, add_new_blocks
5+
from tests.utils import add_blocks_unlock_reward
56

67

78
class DoubleSpendingTestCase(unittest.TestCase):

tests/consensus/test_soft_voided.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22
from hathor.graphviz import GraphvizVisualizer
33
from hathor.simulator import FakeConnection, Simulator
44
from hathor.simulator.trigger import StopAfterNTransactions
5+
from hathor.simulator.utils import gen_new_tx
56
from tests import unittest
67
from tests.simulation.base import SimulatorTestCase
7-
from tests.utils import add_custom_tx, gen_new_tx
8+
from tests.utils import add_custom_tx
89

910
settings = HathorSettings()
1011

tests/consensus/test_soft_voided2.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
from hathor.conf import HathorSettings
22
from hathor.graphviz import GraphvizVisualizer
33
from hathor.simulator import Simulator
4+
from hathor.simulator.utils import gen_new_tx
45
from tests import unittest
56
from tests.simulation.base import SimulatorTestCase
6-
from tests.utils import BURN_ADDRESS, add_custom_tx, gen_new_tx
7+
from tests.utils import BURN_ADDRESS, add_custom_tx
78

89
settings = HathorSettings()
910

tests/consensus/test_soft_voided3.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22
from hathor.graphviz import GraphvizVisualizer
33
from hathor.simulator import FakeConnection, Simulator
44
from hathor.simulator.trigger import StopAfterNTransactions
5+
from hathor.simulator.utils import gen_new_tx
56
from tests import unittest
67
from tests.simulation.base import SimulatorTestCase
7-
from tests.utils import add_custom_tx, gen_custom_tx, gen_new_tx
8+
from tests.utils import add_custom_tx, gen_custom_tx
89

910
settings = HathorSettings()
1011

tests/consensus/test_soft_voided4.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,10 @@
22
from hathor.graphviz import GraphvizVisualizer
33
from hathor.simulator import FakeConnection, Simulator
44
from hathor.simulator.trigger import StopAfterNTransactions
5+
from hathor.simulator.utils import gen_new_double_spending
56
from tests import unittest
67
from tests.simulation.base import SimulatorTestCase
7-
from tests.utils import add_custom_tx, gen_new_double_spending
8+
from tests.utils import add_custom_tx
89

910
settings = HathorSettings()
1011

tests/event/test_event_reorg.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
from hathor.conf import HathorSettings
22
from hathor.event.model.event_type import EventType
33
from hathor.event.storage import EventMemoryStorage
4+
from hathor.simulator.utils import add_new_blocks
45
from tests import unittest
5-
from tests.utils import BURN_ADDRESS, add_new_blocks, get_genesis_key
6+
from tests.utils import BURN_ADDRESS, get_genesis_key
67

78
settings = HathorSettings()
89

0 commit comments

Comments
 (0)