Simplify gen_and_send_sb and make it type safe.

chia/_tests/core/mempool/test_mempool.py

@@ -367,6 +367,18 @@ async def next_block(full_node_1, wallet_a, bt) -> Coin:
 mis = MempoolInclusionStatus
 
 
+async def send_sb(node: FullNodeAPI, sb: SpendBundle) -> Optional[Message]:
+    tx = wallet_protocol.SendTransaction(sb)
+    return await node.send_transaction(tx, test=True)
+
+
+async def gen_and_send_sb(node: FullNodeAPI, wallet: WalletTool, coin: Coin, fee: uint64 = uint64(0)) -> SpendBundle:
+    sb = generate_test_spend_bundle(wallet=wallet, coin=coin, fee=fee)
+    assert sb is not None
+    await send_sb(node, sb)
+    return sb
+
+
 class TestMempoolManager:
     @pytest.mark.anyio
     async def test_basic_mempool_manager(self, two_nodes_one_block, wallet_a, self_hostname):
@@ -548,17 +560,6 @@ async def test_double_spend(self, two_nodes_one_block, wallet_a, self_hostname):
         assert sb2 is None
         assert status == MempoolInclusionStatus.PENDING
 
-    async def send_sb(self, node: FullNodeAPI, sb: SpendBundle) -> Optional[Message]:
-        tx = wallet_protocol.SendTransaction(sb)
-        return await node.send_transaction(tx, test=True)
-
-    async def gen_and_send_sb(self, node, peer, *args, **kwargs):
-        sb = generate_test_spend_bundle(*args, **kwargs)
-        assert sb is not None
-
-        await self.send_sb(node, sb)
-        return sb
-
     def assert_sb_in_pool(self, node, sb):
         assert sb == node.full_node.mempool_manager.get_spendbundle(sb.name())
 
@@ -579,7 +580,7 @@ async def test_double_spend_with_higher_fee(self, two_nodes_one_block, wallet_a,
             farmer_reward_puzzle_hash=reward_ph,
             pool_reward_puzzle_hash=reward_ph,
         )
-        peer = await connect_and_get_peer(server_1, server_2, self_hostname)
+        await connect_and_get_peer(server_1, server_2, self_hostname)
 
         invariant_check_mempool(full_node_1.full_node.mempool_manager.mempool)
         for block in blocks:
@@ -591,15 +592,15 @@ async def test_double_spend_with_higher_fee(self, two_nodes_one_block, wallet_a,
         coins = iter(blocks[-2].get_included_reward_coins())
         coin3, coin4 = next(coins), next(coins)
 
-        sb1_1 = await self.gen_and_send_sb(full_node_1, peer, wallet_a, coin1)
-        sb1_2 = await self.gen_and_send_sb(full_node_1, peer, wallet_a, coin1, fee=uint64(1))
+        sb1_1 = await gen_and_send_sb(full_node_1, wallet_a, coin1)
+        sb1_2 = await gen_and_send_sb(full_node_1, wallet_a, coin1, fee=uint64(1))
 
         # Fee increase is insufficient, the old spendbundle must stay
         self.assert_sb_in_pool(full_node_1, sb1_1)
         self.assert_sb_not_in_pool(full_node_1, sb1_2)
         invariant_check_mempool(full_node_1.full_node.mempool_manager.mempool)
 
-        sb1_3 = await self.gen_and_send_sb(full_node_1, peer, wallet_a, coin1, fee=MEMPOOL_MIN_FEE_INCREASE)
+        sb1_3 = await gen_and_send_sb(full_node_1, wallet_a, coin1, fee=MEMPOOL_MIN_FEE_INCREASE)
 
         # Fee increase is sufficiently high, sb1_1 gets replaced with sb1_3
         self.assert_sb_not_in_pool(full_node_1, sb1_1)
@@ -608,7 +609,7 @@ async def test_double_spend_with_higher_fee(self, two_nodes_one_block, wallet_a,
 
         sb2 = generate_test_spend_bundle(wallet_a, coin2, fee=MEMPOOL_MIN_FEE_INCREASE)
         sb12 = SpendBundle.aggregate((sb2, sb1_3))
-        await self.send_sb(full_node_1, sb12)
+        await send_sb(full_node_1, sb12)
 
         # Aggregated spendbundle sb12 replaces sb1_3 since it spends a superset
         # of coins spent in sb1_3
@@ -618,29 +619,29 @@ async def test_double_spend_with_higher_fee(self, two_nodes_one_block, wallet_a,
 
         sb3 = generate_test_spend_bundle(wallet_a, coin3, fee=uint64(MEMPOOL_MIN_FEE_INCREASE * 2))
         sb23 = SpendBundle.aggregate((sb2, sb3))
-        await self.send_sb(full_node_1, sb23)
+        await send_sb(full_node_1, sb23)
 
         # sb23 must not replace existing sb12 as the former does not spend all
         # coins that are spent in the latter (specifically, coin1)
         self.assert_sb_in_pool(full_node_1, sb12)
         self.assert_sb_not_in_pool(full_node_1, sb23)
         invariant_check_mempool(full_node_1.full_node.mempool_manager.mempool)
 
-        await self.send_sb(full_node_1, sb3)
+        await send_sb(full_node_1, sb3)
         # Adding non-conflicting sb3 should succeed
         self.assert_sb_in_pool(full_node_1, sb3)
         invariant_check_mempool(full_node_1.full_node.mempool_manager.mempool)
 
         sb4_1 = generate_test_spend_bundle(wallet_a, coin4, fee=MEMPOOL_MIN_FEE_INCREASE)
         sb1234_1 = SpendBundle.aggregate((sb12, sb3, sb4_1))
-        await self.send_sb(full_node_1, sb1234_1)
+        await send_sb(full_node_1, sb1234_1)
         # sb1234_1 should not be in pool as it decreases total fees per cost
         self.assert_sb_not_in_pool(full_node_1, sb1234_1)
         invariant_check_mempool(full_node_1.full_node.mempool_manager.mempool)
 
         sb4_2 = generate_test_spend_bundle(wallet_a, coin4, fee=uint64(MEMPOOL_MIN_FEE_INCREASE * 2))
         sb1234_2 = SpendBundle.aggregate((sb12, sb3, sb4_2))
-        await self.send_sb(full_node_1, sb1234_2)
+        await send_sb(full_node_1, sb1234_2)
         # sb1234_2 has a higher fee per cost than its conflicts and should get
         # into mempool
         self.assert_sb_in_pool(full_node_1, sb1234_2)
@@ -673,7 +674,7 @@ async def test_invalid_signature(self, one_node_one_block, wallet_a):
         sb: SpendBundle = generate_test_spend_bundle(wallet_a, coin1)
         assert sb.aggregated_signature != G2Element.generator()
         sb = dataclasses.replace(sb, aggregated_signature=G2Element.generator())
-        res: Optional[Message] = await self.send_sb(full_node_1, sb)
+        res: Optional[Message] = await send_sb(full_node_1, sb)
         assert res is not None
         ack: TransactionAck = TransactionAck.from_bytes(res.data)
         assert ack.status == MempoolInclusionStatus.FAILED.value