Skip to content

Commit 85c5ae3

Browse files
committed
refactor: move simulator utils functions
1 parent 3794dd5 commit 85c5ae3

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

+285
-206
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

+4-3
Original file line numberDiff line numberDiff line change
@@ -44,15 +44,16 @@ 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 import daa
5454
from hathor.conf.get_settings import get_settings
55-
from tests.utils import add_new_blocks, gen_new_tx
55+
from hathor.simulator.utils import add_new_blocks
56+
from hathor.simulator.utils import gen_new_tx
5657

5758
settings = get_settings()
5859
assert manager.wallet is not None
@@ -79,7 +80,7 @@ def simulate_single_chain_blocks_and_transactions(simulator: 'Simulator', manage
7980

8081
def simulate_reorg(simulator: 'Simulator', manager: 'HathorManager') -> None:
8182
from hathor.simulator import FakeConnection
82-
from tests.utils import add_new_blocks
83+
from hathor.simulator.utils import add_new_blocks
8384

8485
builder = simulator.get_default_builder()
8586
manager2 = simulator.create_peer(builder)

hathor/simulator/tx_generator.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121
from hathor.conf.get_settings import get_settings
2222
from hathor.transaction.exceptions import RewardLocked
2323
from hathor.util import Random
24+
from hathor.simulator.utils import gen_new_tx, NoCandidatesError, gen_new_double_spending
2425
from hathor.wallet.exceptions import InsufficientFunds
25-
from tests.utils import NoCandidatesError, gen_new_double_spending, gen_new_tx
2626

2727
if TYPE_CHECKING:
2828
from hathor.manager import HathorManager

hathor/simulator/utils.py

+177
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
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 Transaction, Block
20+
from hathor.types import VertexId, Address
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+
tx = manager.wallet.prepare_transaction_compute_inputs(Transaction, outputs, manager.tx_storage)
42+
tx.storage = manager.tx_storage
43+
44+
max_ts_spent_tx = max(tx.get_spent_tx(txin).timestamp for txin in tx.inputs)
45+
tx.timestamp = max(max_ts_spent_tx + 1, int(manager.reactor.seconds()))
46+
47+
tx.weight = 1
48+
tx.parents = manager.get_new_tx_parents(tx.timestamp)
49+
tx.resolve()
50+
if verify:
51+
tx.verify()
52+
return tx
53+
54+
55+
def add_new_blocks(
56+
manager: HathorManager,
57+
num_blocks: int,
58+
advance_clock: int = None,
59+
*,
60+
parent_block_hash: Optional[VertexId] = None,
61+
block_data: bytes = b'',
62+
weight: Optional[float] = None,
63+
address: Optional[Address] = None
64+
) -> list[Block]:
65+
""" Create, resolve and propagate some blocks
66+
67+
:param manager: Manager object to handle the creation
68+
:type manager: :py:class:`hathor.manager.HathorManager`
69+
70+
:param num_blocks: Quantity of blocks to be created
71+
:type num_blocks: int
72+
73+
:return: Blocks created
74+
:rtype: list[Block]
75+
"""
76+
blocks = []
77+
for _ in range(num_blocks):
78+
blocks.append(
79+
add_new_block(manager, advance_clock, parent_block_hash=parent_block_hash,
80+
data=block_data, weight=weight, address=address)
81+
)
82+
if parent_block_hash:
83+
parent_block_hash = blocks[-1].hash
84+
return blocks
85+
86+
87+
def add_new_block(
88+
manager: HathorManager,
89+
advance_clock: int = None,
90+
*,
91+
parent_block_hash: Optional[VertexId] = None,
92+
data: bytes = b'',
93+
weight: Optional[float] = None,
94+
address: Optional[Address] = None,
95+
propagate: bool = True
96+
) -> Block:
97+
""" Create, resolve and propagate a new block
98+
99+
:param manager: Manager object to handle the creation
100+
:type manager: :py:class:`hathor.manager.HathorManager`
101+
102+
:return: Block created
103+
:rtype: :py:class:`hathor.transaction.block.Block`
104+
"""
105+
block = manager.generate_mining_block(parent_block_hash=parent_block_hash, data=data, address=address)
106+
if weight is not None:
107+
block.weight = weight
108+
block.resolve()
109+
block.validate_full()
110+
if propagate:
111+
manager.propagate_tx(block, fails_silently=False)
112+
if advance_clock:
113+
manager.reactor.advance(advance_clock)
114+
return block
115+
116+
117+
class NoCandidatesError(Exception):
118+
pass
119+
120+
121+
def gen_new_double_spending(manager: HathorManager, *, use_same_parents: bool = False,
122+
tx: Optional[Transaction] = None, weight: float = 1) -> Transaction:
123+
"""
124+
Generate and return a double spending transaction.
125+
126+
Args:
127+
manager: the HathorManager to generate the transaction for
128+
use_same_parents: whether to use the same parents as the original transaction
129+
tx: the original transaction do double spend
130+
weight: the new transaction's weight
131+
132+
Returns: the double spending transaction.
133+
"""
134+
if tx is None:
135+
tx_candidates = manager.get_new_tx_parents()
136+
genesis = manager.tx_storage.get_all_genesis()
137+
genesis_txs = [tx for tx in genesis if not tx.is_block]
138+
# XXX: it isn't possible to double-spend a genesis transaction, thus we remove it from tx_candidates
139+
for genesis_tx in genesis_txs:
140+
if genesis_tx.hash in tx_candidates:
141+
tx_candidates.remove(genesis_tx.hash)
142+
if not tx_candidates:
143+
raise NoCandidatesError()
144+
# assert tx_candidates, 'Must not be empty, otherwise test was wrongly set up'
145+
tx_hash = manager.rng.choice(tx_candidates)
146+
tx = cast(Transaction, manager.tx_storage.get_transaction(tx_hash))
147+
148+
txin = manager.rng.choice(tx.inputs)
149+
150+
from hathor.transaction.scripts import P2PKH, parse_address_script
151+
spent_tx = tx.get_spent_tx(txin)
152+
spent_txout = spent_tx.outputs[txin.index]
153+
p2pkh = parse_address_script(spent_txout.script)
154+
assert isinstance(p2pkh, P2PKH)
155+
156+
from hathor.wallet.base_wallet import WalletInputInfo, WalletOutputInfo
157+
value = spent_txout.value
158+
wallet = manager.wallet
159+
assert wallet is not None
160+
private_key = wallet.get_private_key(p2pkh.address)
161+
inputs = [WalletInputInfo(tx_id=txin.tx_id, index=txin.index, private_key=private_key)]
162+
163+
address = wallet.get_unused_address(mark_as_used=True)
164+
outputs = [WalletOutputInfo(address=decode_address(address), value=int(value), timelock=None)]
165+
166+
tx2 = wallet.prepare_transaction(Transaction, inputs, outputs)
167+
tx2.storage = manager.tx_storage
168+
tx2.weight = weight
169+
tx2.timestamp = max(tx.timestamp + 1, int(manager.reactor.seconds()))
170+
171+
if use_same_parents:
172+
tx2.parents = list(tx.parents)
173+
else:
174+
tx2.parents = manager.get_new_tx_parents(tx2.timestamp)
175+
176+
tx2.resolve()
177+
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
@@ -8,10 +8,11 @@
88
from hathor.crypto.util import decode_address
99
from hathor.transaction import Transaction, TxInput, TxOutput
1010
from hathor.transaction.scripts import create_output_script
11+
from hathor.simulator.utils import add_new_blocks
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
@@ -8,10 +8,10 @@
88
from hathor.conf import HathorSettings
99
from hathor.transaction import Transaction, TransactionMetadata
1010
from hathor.util import json_loadb
11+
from hathor.simulator.utils import add_new_blocks
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
@@ -2,15 +2,9 @@
22

33
from hathor.conf import HathorSettings
44
from hathor.transaction.storage import TransactionMemoryStorage
5+
from hathor.simulator.utils import gen_new_tx, add_new_blocks, add_new_block
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_blocks, add_new_block
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
@@ -3,8 +3,9 @@
33
from hathor.conf import HathorSettings
44
from hathor.event.model.event_type import EventType
55
from hathor.event.storage import EventMemoryStorage
6+
from hathor.simulator.utils import add_new_blocks
67
from tests import unittest
7-
from tests.utils import BURN_ADDRESS, add_new_blocks, get_genesis_key
8+
from tests.utils import BURN_ADDRESS, get_genesis_key
89

910
settings = HathorSettings()
1011

tests/others/test_init_manager.py

+2-7
Original file line numberDiff line numberDiff line change
@@ -4,15 +4,10 @@
44
from hathor.pubsub import PubSubManager
55
from hathor.transaction import BaseTransaction
66
from hathor.transaction.storage import TransactionMemoryStorage
7+
from hathor.simulator.utils import add_new_blocks, add_new_block
78
from tests import unittest
89
from tests.unittest import TestBuilder
9-
from tests.utils import (
10-
add_blocks_unlock_reward,
11-
add_new_block,
12-
add_new_blocks,
13-
add_new_double_spending,
14-
add_new_transactions,
15-
)
10+
from tests.utils import add_blocks_unlock_reward, add_new_double_spending, add_new_transactions
1611

1712
settings = HathorSettings()
1813

0 commit comments

Comments
 (0)