refactor: wallet on_new_tx

glevco · glevco · commit 6fc3b8b3da14 · 2026-01-21T11:18:53.000-03:00
diff --git a/hathor/pubsub.py b/hathor/pubsub.py
@@ -36,6 +36,10 @@
 
 class HathorEvents(Enum):
     """
+        NETWORK_NEW_TX_PROCESSING:
+            Triggered when a new tx/block is received and will begin processing, just before consensus
+            Publishes a tx/block object
+
         NETWORK_NEW_TX_ACCEPTED:
             Triggered when a new tx/block is accepted in the network
             Publishes a tx/block object
@@ -118,6 +122,8 @@ class HathorEvents(Enum):
 
     NETWORK_PEER_DISCONNECTED = 'network:peer_disconnected'
 
+    NETWORK_NEW_TX_PROCESSING = 'network:new_tx_processing'
+
     NETWORK_NEW_TX_ACCEPTED = 'network:new_tx_accepted'
 
     CONSENSUS_TX_UPDATE = 'consensus:tx_update'
diff --git a/hathor/vertex_handler/vertex_handler.py b/hathor/vertex_handler/vertex_handler.py
@@ -28,7 +28,7 @@
 from hathor.feature_activation.feature_service import FeatureService
 from hathor.feature_activation.utils import is_fee_active, is_nano_active
 from hathor.profiler import get_cpu_profiler
-from hathor.pubsub import HathorEvents, PubSubManager
+from hathor.pubsub import EventArguments, HathorEvents, PubSubManager
 from hathor.reactor import ReactorProtocol
 from hathor.transaction import BaseTransaction, Block, Transaction
 from hathor.transaction.storage import TransactionStorage
@@ -52,7 +52,6 @@ class VertexHandler:
         '_feature_service',
         '_pubsub',
         '_execution_manager',
-        '_wallet',
         '_log_vertex_bytes',
     )
 
@@ -79,9 +78,21 @@ def __init__(
         self._feature_service = feature_service
         self._pubsub = pubsub
         self._execution_manager = execution_manager
-        self._wallet = wallet
         self._log_vertex_bytes = log_vertex_bytes
 
+        self._register_wallet(wallet)
+
+    def _register_wallet(self, wallet: BaseWallet | None) -> None:
+        """Register a wallet on pubsub."""
+        if wallet is None:
+            return
+
+        def handler(event: HathorEvents, args: EventArguments) -> None:
+            assert event == HathorEvents.NETWORK_NEW_TX_PROCESSING
+            wallet.on_new_tx(args.tx)
+
+        self._pubsub.subscribe(HathorEvents.NETWORK_NEW_TX_PROCESSING, handler)
+
     @cpu.profiler('on_new_block')
     @inlineCallbacks
     def on_new_block(self, block: Block, *, deps: list[Transaction]) -> Generator[Any, Any, bool]:
@@ -239,6 +250,7 @@ def _unsafe_save_and_run_consensus(self, vertex: BaseTransaction) -> None:
         self._tx_storage.save_transaction(vertex)
         with non_critical_code(self._log):
             self._tx_storage.indexes.add_to_non_critical_indexes(vertex)
+        self._pubsub.publish(HathorEvents.NETWORK_NEW_TX_PROCESSING, tx=vertex)
         self._consensus.unsafe_update(vertex)
 
     def _post_consensus(
@@ -268,10 +280,6 @@ def _post_consensus(
             # Publish to pubsub manager the new tx accepted, now that it's full validated
             self._pubsub.publish(HathorEvents.NETWORK_NEW_TX_ACCEPTED, tx=vertex)
 
-            if self._wallet:
-                # TODO Remove it and use pubsub instead.
-                self._wallet.on_new_tx(vertex)
-
             self._log_new_object(vertex, 'new {}', quiet=quiet)
 
     def _log_new_object(self, tx: BaseTransaction, message_fmt: str, *, quiet: bool) -> None:
diff --git a/hathor_cli/events_simulator/scenario.py b/hathor_cli/events_simulator/scenario.py
@@ -177,7 +177,7 @@ def simulate_invalid_mempool_transaction(simulator: 'Simulator', manager: 'Hatho
     assert manager.wallet is not None
     address = manager.wallet.get_unused_address(mark_as_used=False)
 
-    blocks = add_new_blocks(manager, settings.REWARD_SPEND_MIN_BLOCKS + 1)
+    blocks = add_new_blocks(manager, settings.REWARD_SPEND_MIN_BLOCKS + 1, advance_clock=1)
     simulator.run(60)
 
     balance_per_address = manager.wallet.get_balance_per_address(settings.HATHOR_TOKEN_UID)
diff --git a/hathor_tests/consensus/test_soft_voided.py b/hathor_tests/consensus/test_soft_voided.py
@@ -84,6 +84,7 @@ def _run_test(
         assert manager2.wallet is not None
         address = manager2.wallet.get_unused_address(mark_as_used=False)
         value = 1
+        simulator.run_to_completion()
         txC = gen_new_tx(manager2, address, value)
         txC.parents[0] = txA.hash
         txC.timestamp = max(txC.timestamp, txA.timestamp + 1)
diff --git a/hathor_tests/p2p/test_double_spending.py b/hathor_tests/p2p/test_double_spending.py
@@ -24,7 +24,7 @@ def _add_new_transactions(self, manager: HathorManager, num_txs: int) -> list[Tr
         for _ in range(num_txs):
             address = not_none(self.get_address(0))
             value = self.rng.choice([5, 10, 15, 20])
-            tx = add_new_tx(manager, address, value)
+            tx = add_new_tx(manager, address, value, advance_clock=1)
             txs.append(tx)
         return txs
 
diff --git a/hathor_tests/resources/transaction/test_block_at_height.py b/hathor_tests/resources/transaction/test_block_at_height.py
@@ -20,14 +20,14 @@ def test_include_full(self):
 
         confirmed_tx_list = []
         for _ in range(15):
-            confirmed_tx_list.append(add_new_tx(self.manager, address, 1))
+            confirmed_tx_list.append(add_new_tx(self.manager, address, 1, advance_clock=1))
 
         block = add_new_block(self.manager, advance_clock=1)
         height = block.get_height()
 
         # non-confirmed transactions
         for _ in range(15):
-            add_new_tx(self.manager, address, 1)
+            add_new_tx(self.manager, address, 1, advance_clock=1)
 
         response = yield self.web.get("block_at_height", {
             b'height': str(height).encode('ascii'),
@@ -48,14 +48,14 @@ def test_include_txids(self):
 
         confirmed_tx_list = []
         for _ in range(15):
-            confirmed_tx_list.append(add_new_tx(self.manager, address, 1))
+            confirmed_tx_list.append(add_new_tx(self.manager, address, 1, advance_clock=1))
 
         block = add_new_block(self.manager, advance_clock=1)
         height = block.get_height()
 
         # non-confirmed transactions
         for _ in range(15):
-            add_new_tx(self.manager, address, 1)
+            add_new_tx(self.manager, address, 1, advance_clock=1)
 
         response = yield self.web.get("block_at_height", {
             b'height': str(height).encode('ascii'),
diff --git a/hathor_tests/resources/transaction/test_transaction_confirmation.py b/hathor_tests/resources/transaction/test_transaction_confirmation.py
@@ -27,7 +27,7 @@ def test_get_data(self):
         # Adding blocks to have funds
         add_new_blocks(self.manager, 2, advance_clock=1)
         add_blocks_unlock_reward(self.manager)
-        tx = add_new_transactions(self.manager, 5)[0]
+        tx = add_new_transactions(self.manager, 5, advance_clock=1)[0]
         add_new_blocks(self.manager, 2, advance_clock=1)
         add_blocks_unlock_reward(self.manager)
         response_success2 = yield self.web.get(
diff --git a/hathor_tests/resources/transaction/test_tx.py b/hathor_tests/resources/transaction/test_tx.py
@@ -292,7 +292,9 @@ def test_get_many(self):
         # Add some blocks and txs and get them in timestamp order
         blocks = add_new_blocks(self.manager, 4, advance_clock=1)
         _blocks = add_blocks_unlock_reward(self.manager)
-        txs = sorted(add_new_transactions(self.manager, 25), key=lambda x: (x.timestamp, x.hash))
+        txs = sorted(
+            add_new_transactions(self.manager, 25, advance_clock=1), key=lambda x: (x.timestamp, x.hash)
+        )
 
         blocks.extend(_blocks)
         blocks = sorted(blocks, key=lambda x: (x.timestamp, x.hash))
@@ -484,7 +486,7 @@ def test_negative_timestamp(self):
         # Add some blocks and txs and get them in timestamp order
         blocks = add_new_blocks(self.manager, 4, advance_clock=1)
         add_blocks_unlock_reward(self.manager)
-        add_new_transactions(self.manager, 25)
+        add_new_transactions(self.manager, 25, advance_clock=1)
 
         response = yield self.web.get(
                 "transaction", {
diff --git a/hathor_tests/resources/wallet/test_balance.py b/hathor_tests/resources/wallet/test_balance.py
@@ -30,6 +30,7 @@ def test_get(self):
             cpu_mining_service=CpuMiningService()
         )
         yield self.web_mining.post("mining", {'block_bytes': base64.b64encode(block_bytes).decode('utf-8')})
+        self.clock.advance(1)
 
         # Get new balance after block
         response2 = yield self.web.get("wallet/balance")
diff --git a/hathor_tests/resources/wallet/test_history.py b/hathor_tests/resources/wallet/test_history.py
@@ -25,6 +25,7 @@ def test_get(self):
             cpu_mining_service=CpuMiningService()
         )
         yield self.web_mining.post("mining", {'block_bytes': base64.b64encode(block_bytes).decode('utf-8')})
+        self.clock.advance(1)
 
         # Getting wallet history
         response = yield self.web.get("wallet/history", {b'page': 1, b'count': 10})
diff --git a/hathor_tests/tx/test_indexes.py b/hathor_tests/tx/test_indexes.py
@@ -362,7 +362,7 @@ def test_utxo_index_limits(self):
         txs = []
         values = list(range(1, 301))
         for value in values:
-            txs.append(add_new_tx(self.manager, address, value))
+            txs.append(add_new_tx(self.manager, address, value, advance_clock=1))
         assert len(txs) == len(values)
         txs_and_values = list(zip(txs, values))
 
diff --git a/hathor_tests/tx/test_indexes4.py b/hathor_tests/tx/test_indexes4.py
@@ -49,6 +49,7 @@ def _build_randomized_blockchain(self, *, utxo_index=False):
             value = 500
             tx = gen_new_tx(manager, address, value)
             assert manager.propagate_tx(tx)
+            self.clock.advance(1)
         return manager
 
     def test_index_initialization(self):
diff --git a/hathor_tests/tx/test_multisig.py b/hathor_tests/tx/test_multisig.py
@@ -123,6 +123,7 @@ def test_spend_multisig(self):
 
         # Now we propagate the correct
         self.assertTrue(self.manager.propagate_tx(tx))
+        self.clock.advance(1)
 
         self.assertEqual(self.manager.wallet.balance[self._settings.HATHOR_TOKEN_UID],
                          WalletBalance(0, first_block_amount + 300))
diff --git a/hathor_tests/tx/test_reward_lock.py b/hathor_tests/tx/test_reward_lock.py
@@ -165,6 +165,7 @@ def test_mempool_tx_invalid_after_reorg(self) -> None:
         assert tx_address not in balance_per_address
         self.assertEqual(tx.static_metadata.min_height, unlock_height)
         self.assertTrue(self.manager.on_new_tx(tx))
+        self.clock.advance(1)
         balance_per_address = self.manager.wallet.get_balance_per_address(self._settings.HATHOR_TOKEN_UID)
         assert balance_per_address[tx_address] == 6400
 
diff --git a/hathor_tests/tx/test_timelock.py b/hathor_tests/tx/test_timelock.py
@@ -40,6 +40,7 @@ def test_timelock(self):
         tx1.timestamp = int(self.clock.seconds())
         self.manager.cpu_mining_service.resolve(tx1)
         self.manager.propagate_tx(tx1)
+        self.clock.advance(1)
 
         self.assertEqual(self.manager.wallet.balance[self._settings.HATHOR_TOKEN_UID],
                          WalletBalance(500, sum(blocks_tokens) - 500))
@@ -79,6 +80,7 @@ def test_timelock(self):
         tx3.timestamp = int(self.clock.seconds())
         self.manager.cpu_mining_service.resolve(tx3)
         propagated = self.manager.propagate_tx(tx3)
+        self.clock.advance(1)
         self.assertEqual(self.manager.wallet.balance[self._settings.HATHOR_TOKEN_UID],
                          WalletBalance(500, sum(blocks_tokens) - 500 - 700))
         self.assertTrue(propagated)
@@ -99,6 +101,7 @@ def test_timelock(self):
         tx4.timestamp = int(self.clock.seconds())
         self.manager.cpu_mining_service.resolve(tx4)
         propagated = self.manager.propagate_tx(tx4)
+        self.clock.advance(1)
         self.assertEqual(self.manager.wallet.balance[self._settings.HATHOR_TOKEN_UID],
                          WalletBalance(500, sum(blocks_tokens[:3])))
         self.assertTrue(propagated)
@@ -107,6 +110,7 @@ def test_timelock(self):
         tx2.timestamp = int(self.clock.seconds())
         self.manager.cpu_mining_service.resolve(tx2)
         propagated = self.manager.propagate_tx(tx2)
+        self.clock.advance(1)
         self.assertEqual(self.manager.wallet.balance[self._settings.HATHOR_TOKEN_UID],
                          WalletBalance(0, sum(blocks_tokens[:3])))
         self.assertTrue(propagated)
diff --git a/hathor_tests/tx/test_traversal.py b/hathor_tests/tx/test_traversal.py
@@ -20,7 +20,7 @@ def setUp(self):
 
         self.blocks_before = add_new_blocks(self.manager, 3, advance_clock=1)
         self.blocks_before.extend(add_blocks_unlock_reward(self.manager))
-        self.txs_before = add_new_transactions(self.manager, 5)
+        self.txs_before = add_new_transactions(self.manager, 5, advance_clock=1)
         for block in self.blocks_before:
             self.hashes_before.add(block.hash)
         for tx in self.txs_before:
@@ -30,7 +30,7 @@ def setUp(self):
         self.root_tx = add_new_tx(self.manager, address=address, value=100)
 
         self.blocks_after = add_blocks_unlock_reward(self.manager)
-        self.txs_after = add_new_transactions(self.manager, 5)
+        self.txs_after = add_new_transactions(self.manager, 5, advance_clock=1)
         self.blocks_after.extend(add_new_blocks(self.manager, 3, advance_clock=1))
 
         self.hashes_after = set()
diff --git a/hathor_tests/wallet/test_wallet_hd.py b/hathor_tests/wallet/test_wallet_hd.py
@@ -26,7 +26,7 @@ def test_transaction_and_balance(self):
         # generate a new block and check if we increase balance
         new_address = self.wallet.get_unused_address()
         out = WalletOutputInfo(decode_address(new_address), self.TOKENS, timelock=None)
-        block = add_new_block(self.manager)
+        block = add_new_block(self.manager, advance_clock=1)
         self.manager.verification_service.verify(block, self.get_verification_params(self.manager))
         utxo = self.wallet.unspent_txs[self._settings.HATHOR_TOKEN_UID].get((block.hash, 0))
         self.assertIsNotNone(utxo)

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,7 @@ def test_get(self):`
`30`	`30`	`cpu_mining_service=CpuMiningService()`
`31`	`31`	`)`
`32`	`32`	`yield self.web_mining.post("mining", {'block_bytes': base64.b64encode(block_bytes).decode('utf-8')})`
	`33`	`+ self.clock.advance(1)`
`33`	`34`
`34`	`35`	`# Get new balance after block`
`35`	`36`	`response2 = yield self.web.get("wallet/balance")`
Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,7 @@ def test_get(self):`
`25`	`25`	`cpu_mining_service=CpuMiningService()`
`26`	`26`	`)`
`27`	`27`	`yield self.web_mining.post("mining", {'block_bytes': base64.b64encode(block_bytes).decode('utf-8')})`
	`28`	`+ self.clock.advance(1)`
`28`	`29`
`29`	`30`	`# Getting wallet history`
`30`	`31`	`response = yield self.web.get("wallet/history", {b'page': 1, b'count': 10})`