Merge pull request #3020 from opentensor/fix/thewhaleking/broken-e2e-tests

Fixes broken e2e tests

Author: basfroman

bittensor/core/async_subtensor.py

@@ -5548,7 +5548,7 @@ async def unstake(
         self,
         wallet: "Wallet",
         hotkey_ss58: Optional[str] = None,
-        netuid: Optional[int] = None,
+        netuid: Optional[int] = None,  # TODO why is this optional?
         amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,

bittensor/core/extrinsics/asyncex/registration.py

@@ -10,6 +10,7 @@
 import asyncio
 from typing import Optional, Union, TYPE_CHECKING
 
+from bittensor.core.extrinsics.asyncex.utils import get_extrinsic_fee
 from bittensor.utils import unlock_key
 from bittensor.utils.btlogging import logging
 from bittensor.utils.registration import log_no_torch_error, create_pow_async, torch
@@ -57,6 +58,12 @@ async def _do_burned_register(
             "hotkey": wallet.hotkey.ss58_address,
         },
     )
+    fee = await get_extrinsic_fee(
+        subtensor=subtensor, call=call, keypair=wallet.coldkeypub
+    )
+    logging.info(
+        f"The registration fee for SN #[blue]{netuid}[/blue] is [blue]{fee}[/blue]."
+    )
     return await subtensor.sign_and_send_extrinsic(
         call=call,
         wallet=wallet,
@@ -127,7 +134,6 @@ async def burned_register_extrinsic(
         return True
 
     logging.debug(":satellite: [magenta]Recycling TAO for Registration...[/magenta]")
-    logging.info(f"Recycling {recycle_amount} to register on subnet:{netuid}")
 
     success, err_msg = await _do_burned_register(
         subtensor=subtensor,

bittensor/core/extrinsics/asyncex/unstaking.py

@@ -2,7 +2,7 @@
 from typing import Optional, TYPE_CHECKING
 
 from async_substrate_interface.errors import SubstrateRequestException
-
+from bittensor.core.extrinsics.asyncex.utils import get_extrinsic_fee
 from bittensor.core.extrinsics.utils import get_old_stakes
 from bittensor.utils import unlock_key, format_error_message
 from bittensor.utils.balance import Balance
@@ -114,14 +114,14 @@ async def unstake_extrinsic(
             else:
                 price_with_tolerance = base_price * (1 - rate_tolerance)
 
-            logging.info(
+            logging_info = (
                 f":satellite: [magenta]Safe Unstaking from:[/magenta] "
                 f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green], "
                 f"tolerance percentage: [green]{rate_tolerance * 100}%[/green], "
                 f"price limit: [green]{price_with_tolerance}[/green], "
                 f"original price: [green]{base_price}[/green], "
                 f"with partial unstake: [green]{allow_partial_stake}[/green] "
-                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+                f"on [blue]{subtensor.network}[/blue]"
             )
 
             limit_price = Balance.from_tao(price_with_tolerance).rao
@@ -133,10 +133,10 @@ async def unstake_extrinsic(
             )
             call_function = "remove_stake_limit"
         else:
-            logging.info(
+            logging_info = (
                 f":satellite: [magenta]Unstaking from:[/magenta] "
                 f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green] "
-                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+                f"on [blue]{subtensor.network}[/blue]"
             )
             call_function = "remove_stake"
 
@@ -145,6 +145,10 @@ async def unstake_extrinsic(
             call_function=call_function,
             call_params=call_params,
         )
+        fee = await get_extrinsic_fee(
+            subtensor=subtensor, call=call, keypair=wallet.coldkeypub, netuid=netuid
+        )
+        logging.info(f"{logging_info} for fee [blue]{fee}[/blue][magenta]...[/magenta]")
         success, message = await subtensor.sign_and_send_extrinsic(
             call=call,
             wallet=wallet,
@@ -381,10 +385,6 @@ async def unstake_multiple_extrinsic(
             continue
 
         try:
-            logging.info(
-                f"Unstaking [blue]{unstaking_balance}[/blue] from hotkey: [magenta]{hotkey_ss58}[/magenta] on netuid: "
-                f"[blue]{netuid}[/blue]"
-            )
             call = await subtensor.substrate.compose_call(
                 call_module="SubtensorModule",
                 call_function="remove_stake",
@@ -394,6 +394,13 @@ async def unstake_multiple_extrinsic(
                     "netuid": netuid,
                 },
             )
+            fee = await get_extrinsic_fee(
+                subtensor=subtensor, call=call, keypair=wallet.coldkeypub, netuid=netuid
+            )
+            logging.info(
+                f"Unstaking [blue]{unstaking_balance}[/blue] from hotkey: [magenta]{hotkey_ss58}[/magenta] on netuid: "
+                f"[blue]{netuid}[/blue] for fee [blue]{fee}[/blue]"
+            )
 
             staking_response, err_msg = await subtensor.sign_and_send_extrinsic(
                 call=call,

bittensor/core/extrinsics/asyncex/utils.py

@@ -0,0 +1,34 @@
+from typing import TYPE_CHECKING, Optional
+
+from bittensor.utils.balance import Balance
+
+if TYPE_CHECKING:
+    from scalecodec import GenericCall
+    from bittensor_wallet import Keypair
+    from bittensor.core.async_subtensor import AsyncSubtensor
+
+
+async def get_extrinsic_fee(
+    subtensor: "AsyncSubtensor",
+    call: "GenericCall",
+    keypair: "Keypair",
+    netuid: Optional[int] = None,
+):
+    """
+    Get extrinsic fee for a given extrinsic call and keypair for a given SN's netuid.
+
+    Arguments:
+        subtensor: The Subtensor instance.
+        netuid: The SN's netuid.
+        call: The extrinsic call.
+        keypair: The keypair associated with the extrinsic.
+
+    Returns:
+        Balance object representing the extrinsic fee in RAO.
+    """
+    payment_info = await subtensor.substrate.get_payment_info(
+        call=call, keypair=keypair
+    )
+    return Balance.from_rao(amount=payment_info["partial_fee"]).set_unit(
+        netuid=netuid or 0
+    )

bittensor/core/extrinsics/registration.py

@@ -9,6 +9,7 @@
 import time
 from typing import Optional, Union, TYPE_CHECKING
 
+from bittensor.core.extrinsics.utils import get_extrinsic_fee
 from bittensor.utils import unlock_key
 from bittensor.utils.btlogging import logging
 from bittensor.utils.registration import create_pow, log_no_torch_error, torch
@@ -56,6 +57,10 @@ def _do_burned_register(
             "hotkey": wallet.hotkey.ss58_address,
         },
     )
+    fee = get_extrinsic_fee(subtensor=subtensor, call=call, keypair=wallet.coldkeypub)
+    logging.info(
+        f"The registration fee for SN #[blue]{netuid}[/blue] is [blue]{fee}[/blue]."
+    )
     return subtensor.sign_and_send_extrinsic(
         call=call,
         wallet=wallet,
@@ -197,6 +202,7 @@ def _do_pow_register(
             "coldkey": wallet.coldkeypub.ss58_address,
         },
     )
+    logging.debug(":satellite: [magenta]Sending POW Register Extrinsic...[/magenta]")
     return subtensor.sign_and_send_extrinsic(
         call=call,
         wallet=wallet,

bittensor/core/extrinsics/unstaking.py

@@ -1,6 +1,7 @@
 from typing import Optional, TYPE_CHECKING
 
 from async_substrate_interface.errors import SubstrateRequestException
+from bittensor.core.extrinsics.utils import get_extrinsic_fee
 
 from bittensor.core.extrinsics.utils import get_old_stakes
 from bittensor.utils import unlock_key, format_error_message
@@ -112,14 +113,14 @@ def unstake_extrinsic(
             else:
                 price_with_tolerance = base_price * (1 - rate_tolerance)
 
-            logging.info(
+            logging_info = (
                 f":satellite: [magenta]Safe Unstaking from:[/magenta] "
                 f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green], "
                 f"tolerance percentage: [green]{rate_tolerance * 100}%[/green], "
                 f"price limit: [green]{price_with_tolerance}[/green], "
                 f"original price: [green]{base_price}[/green], "
                 f"with partial unstake: [green]{allow_partial_stake}[/green] "
-                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+                f"on [blue]{subtensor.network}[/blue]"
             )
 
             limit_price = Balance.from_tao(price_with_tolerance).rao
@@ -131,10 +132,10 @@ def unstake_extrinsic(
             )
             call_function = "remove_stake_limit"
         else:
-            logging.info(
+            logging_info = (
                 f":satellite: [magenta]Unstaking from:[/magenta] "
                 f"netuid: [green]{netuid}[/green], amount: [green]{unstaking_balance}[/green] "
-                f"on [blue]{subtensor.network}[/blue][magenta]...[/magenta]"
+                f"on [blue]{subtensor.network}[/blue]"
             )
             call_function = "remove_stake"
 
@@ -143,6 +144,10 @@ def unstake_extrinsic(
             call_function=call_function,
             call_params=call_params,
         )
+        fee = get_extrinsic_fee(
+            subtensor=subtensor, netuid=netuid, call=call, keypair=wallet.coldkeypub
+        )
+        logging.info(f"{logging_info} for fee [blue]{fee}[/blue][magenta]...[/magenta]")
 
         success, message = subtensor.sign_and_send_extrinsic(
             call=call,
@@ -372,9 +377,6 @@ def unstake_multiple_extrinsic(
             continue
 
         try:
-            logging.info(
-                f"Unstaking [blue]{unstaking_balance}[/blue] from [magenta]{hotkey_ss58}[/magenta] on [blue]{netuid}[/blue]"
-            )
             call = subtensor.substrate.compose_call(
                 call_module="SubtensorModule",
                 call_function="remove_stake",
@@ -384,6 +386,14 @@ def unstake_multiple_extrinsic(
                     "netuid": netuid,
                 },
             )
+            fee = get_extrinsic_fee(
+                subtensor=subtensor, netuid=netuid, call=call, keypair=wallet.coldkeypub
+            )
+            logging.info(
+                f"Unstaking [blue]{unstaking_balance}[/blue] from hotkey: [magenta]{hotkey_ss58}[/magenta] on netuid: "
+                f"[blue]{netuid}[/blue] for fee [blue]{fee}[/blue]"
+            )
+
             staking_response, err_msg = subtensor.sign_and_send_extrinsic(
                 call=call,
                 wallet=wallet,

bittensor/core/extrinsics/utils.py

@@ -1,12 +1,14 @@
 """Module with helper functions for extrinsics."""
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 from bittensor.utils.balance import Balance
 
 if TYPE_CHECKING:
-    from bittensor_wallet import Wallet
+    from scalecodec import GenericCall
+    from bittensor_wallet import Wallet, Keypair
     from bittensor.core.chain_data import StakeInfo
+    from bittensor.core.subtensor import Subtensor
 
 
 def get_old_stakes(
@@ -42,3 +44,27 @@ def get_old_stakes(
         )
         for hotkey_ss58, netuid in zip(hotkey_ss58s, netuids)
     ]
+
+
+def get_extrinsic_fee(
+    call: "GenericCall",
+    keypair: "Keypair",
+    subtensor: "Subtensor",
+    netuid: Optional[int] = None,
+):
+    """
+    Get extrinsic fee for a given extrinsic call and keypair for a given SN's netuid.
+
+    Arguments:
+        subtensor: The Subtensor instance.
+        call: The extrinsic call.
+        keypair: The keypair associated with the extrinsic.
+        netuid: The SN's netuid.
+
+    Returns:
+        Balance object representing the extrinsic fee in RAO.
+    """
+    payment_info = subtensor.substrate.get_payment_info(call=call, keypair=keypair)
+    return Balance.from_rao(amount=payment_info["partial_fee"]).set_unit(
+        netuid=netuid or 0
+    )

bittensor/core/subtensor.py

@@ -4368,7 +4368,7 @@ def unstake(
         self,
         wallet: "Wallet",
         hotkey_ss58: Optional[str] = None,
-        netuid: Optional[int] = None,
+        netuid: Optional[int] = None,  # TODO why is this optional?
         amount: Optional[Balance] = None,
         wait_for_inclusion: bool = True,
         wait_for_finalization: bool = False,