Added e2e tests

Author: thewhaleking

bittensor/core/async_subtensor.py

@@ -3134,7 +3134,7 @@ async def get_total_subnets(
         return getattr(result, "value", None)
 
     async def get_transfer_fee(
-        self, wallet: "Wallet", dest: str, value: Balance
+        self, wallet: "Wallet", dest: str, value: Balance, keep_alive: bool = True
     ) -> Balance:
         """
         Calculates the transaction fee for transferring tokens from a wallet to a specified destination address. This
@@ -3146,6 +3146,8 @@ async def get_transfer_fee(
             dest: The ``SS58`` address of the destination account.
             value: The amount of tokens to be transferred, specified as a Balance object, or in Tao (float) or Rao
                 (int) units.
+            keep_alive: Whether the transfer fee should be calculated based on keeping the wallet alive (existential
+                deposit) or not.
 
         Returns:
             bittensor.utils.balance.Balance: The estimated transaction fee for the transfer, represented as a Balance
@@ -3155,12 +3157,25 @@ async def get_transfer_fee(
         wallet has sufficient funds to cover both the transfer amount and the associated costs. This function provides
         a crucial tool for managing financial operations within the Bittensor network.
         """
-        value = check_and_convert_to_balance(value)
+        call_params = {"dest": dest}
+        if value is None:
+            call_function = "transfer_all"
+            if keep_alive:
+                call_params["keep_alive"] = True
+            else:
+                call_params["keep_alive"] = False
+        else:
+            value = check_and_convert_to_balance(value)
+            call_params["value"] = value.rao
+            if keep_alive:
+                call_function = "transfer_keep_alive"
+            else:
+                call_function = "transfer_allow_death"
 
         call = await self.substrate.compose_call(
             call_module="Balances",
-            call_function="transfer_keep_alive",
-            call_params={"dest": dest, "value": value.rao},
+            call_function=call_function,
+            call_params=call_params,
         )
 
         try:

bittensor/core/extrinsics/asyncex/transfer.py

@@ -150,7 +150,7 @@ async def transfer_extrinsic(
     )
 
     fee = await subtensor.get_transfer_fee(
-        wallet=wallet, dest=destination, value=amount
+        wallet=wallet, dest=destination, value=amount, keep_alive=keep_alive
     )
 
     if not keep_alive:

bittensor/core/extrinsics/transfer.py

@@ -147,7 +147,9 @@ def transfer_extrinsic(
     else:
         existential_deposit = subtensor.get_existential_deposit(block=block)
 
-    fee = subtensor.get_transfer_fee(wallet=wallet, dest=dest, value=amount)
+    fee = subtensor.get_transfer_fee(
+        wallet=wallet, dest=dest, value=amount, keep_alive=keep_alive
+    )
 
     # Check if we have enough balance.
     if transfer_all is True:

bittensor/core/subtensor.py

@@ -2243,7 +2243,13 @@ def get_total_subnets(self, block: Optional[int] = None) -> Optional[int]:
         )
         return getattr(result, "value", None)
 
-    def get_transfer_fee(self, wallet: "Wallet", dest: str, value: Balance) -> Balance:
+    def get_transfer_fee(
+        self,
+        wallet: "Wallet",
+        dest: str,
+        value: Optional[Balance],
+        keep_alive: bool = True,
+    ) -> Balance:
         """
         Calculates the transaction fee for transferring tokens from a wallet to a specified destination address. This
             function simulates the transfer to estimate the associated cost, taking into account the current network
@@ -2254,6 +2260,8 @@ def get_transfer_fee(self, wallet: "Wallet", dest: str, value: Balance) -> Balan
             dest (str): The ``SS58`` address of the destination account.
             value (Union[bittensor.utils.balance.Balance, float, int]): The amount of tokens to be transferred,
                 specified as a Balance object, or in Tao (float) or Rao (int) units.
+            keep_alive: Whether the transfer fee should be calculated based on keeping the wallet alive (existential
+                deposit) or not.
 
         Returns:
             bittensor.utils.balance.Balance: The estimated transaction fee for the transfer, represented as a Balance
@@ -2263,11 +2271,24 @@ def get_transfer_fee(self, wallet: "Wallet", dest: str, value: Balance) -> Balan
             has sufficient funds to cover both the transfer amount and the associated costs. This function provides a
             crucial tool for managing financial operations within the Bittensor network.
         """
-        value = check_and_convert_to_balance(value)
+        call_params = {"dest": dest}
+        if value is None:
+            call_function = "transfer_all"
+            if keep_alive:
+                call_params["keep_alive"] = True
+            else:
+                call_params["keep_alive"] = False
+        else:
+            value = check_and_convert_to_balance(value)
+            call_params["value"] = value.rao
+            if keep_alive:
+                call_function = "transfer_keep_alive"
+            else:
+                call_function = "transfer_allow_death"
         call = self.substrate.compose_call(
             call_module="Balances",
-            call_function="transfer_keep_alive",
-            call_params={"dest": dest, "value": value.rao},
+            call_function=call_function,
+            call_params=call_params,
         )
 
         try:

tests/e2e_tests/test_transfer.py

@@ -1,10 +1,18 @@
+import typing
+
+from bittensor_wallet import Wallet
+import pytest
+
 from bittensor.utils.balance import Balance
 from bittensor import logging
 
+if typing.TYPE_CHECKING:
+    from bittensor.core.subtensor_api import SubtensorApi
+
 logging.set_trace()
 
 
-def test_transfer(subtensor, alice_wallet):
+def test_transfer(subtensor: "SubtensorApi", alice_wallet):
     """
     Test the transfer mechanism on the chain
 
@@ -47,3 +55,90 @@ def test_transfer(subtensor, alice_wallet):
     )
 
     print("✅ Passed test_transfer")
+
+
+def test_transfer_all(subtensor: "Subtensor", alice_wallet):
+    # create two dummy accounts we can drain
+    dummy_account_1 = Wallet(path="/tmp/bittensor-dummy-account-1")
+    dummy_account_2 = Wallet(path="/tmp/bittensor-dummy-account-2")
+    dummy_account_1.create_new_coldkey(use_password=False, overwrite=True)
+    dummy_account_2.create_new_coldkey(use_password=False, overwrite=True)
+
+    # fund the first dummy account
+    assert subtensor.transfer(
+        alice_wallet,
+        dest=dummy_account_1.coldkeypub.ss58_address,
+        amount=Balance.from_tao(2.0),
+        wait_for_finalization=True,
+        wait_for_inclusion=True,
+    )
+    # Account details before transfer
+    existential_deposit = subtensor.get_existential_deposit()
+    assert subtensor.transfer(
+        wallet=dummy_account_1,
+        dest=dummy_account_2.coldkeypub.ss58_address,
+        amount=None,
+        transfer_all=True,
+        wait_for_finalization=True,
+        wait_for_inclusion=True,
+        keep_alive=True,
+    )
+    balance_after = subtensor.get_balance(dummy_account_1.coldkeypub.ss58_address)
+    assert balance_after == existential_deposit
+    assert subtensor.transfer(
+        wallet=dummy_account_2,
+        dest=alice_wallet.coldkeypub.ss58_address,
+        amount=None,
+        transfer_all=True,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        keep_alive=False,
+    )
+    balance_after = subtensor.get_balance(dummy_account_2.coldkeypub.ss58_address)
+    assert balance_after == Balance(0)
+
+
+@pytest.mark.asyncio
+async def test_async_transfer(async_subtensor: "SubtensorApi", alice_wallet):
+    # create two dummy accounts we can drain
+    dummy_account_1 = Wallet(path="/tmp/bittensor-dummy-account-3")
+    dummy_account_2 = Wallet(path="/tmp/bittensor-dummy-account-4")
+    dummy_account_1.create_new_coldkey(use_password=False, overwrite=True)
+    dummy_account_2.create_new_coldkey(use_password=False, overwrite=True)
+
+    # fund the first dummy account
+    assert await async_subtensor.transfer(
+        alice_wallet,
+        dest=dummy_account_1.coldkeypub.ss58_address,
+        amount=Balance.from_tao(2.0),
+        wait_for_finalization=True,
+        wait_for_inclusion=True,
+    )
+    # Account details before transfer
+    existential_deposit = await async_subtensor.get_existential_deposit()
+    assert await async_subtensor.transfer(
+        wallet=dummy_account_1,
+        dest=dummy_account_2.coldkeypub.ss58_address,
+        amount=None,
+        transfer_all=True,
+        wait_for_finalization=True,
+        wait_for_inclusion=True,
+        keep_alive=True,
+    )
+    balance_after = await async_subtensor.get_balance(
+        dummy_account_1.coldkeypub.ss58_address
+    )
+    assert balance_after == existential_deposit
+    assert await async_subtensor.transfer(
+        wallet=dummy_account_2,
+        dest=alice_wallet.coldkeypub.ss58_address,
+        amount=None,
+        transfer_all=True,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        keep_alive=False,
+    )
+    balance_after = await async_subtensor.get_balance(
+        dummy_account_2.coldkeypub.ss58_address
+    )
+    assert balance_after == Balance(0)