Adds e2e test for safe add,remove stake

Author: ibraheem-abe

tests/e2e_tests/test_staking.py

@@ -3,6 +3,9 @@
 from bittensor.core.chain_data.stake_info import StakeInfo
 from bittensor.utils.balance import Balance
 from tests.e2e_tests.utils.chain_interactions import ANY_BALANCE
+from bittensor import logging
+
+logging.enable_info()
 
 
 def test_single_operation(subtensor, alice_wallet, bob_wallet):
@@ -245,3 +248,169 @@ def test_batch_operations(subtensor, alice_wallet, bob_wallet):
         bob_wallet.coldkey.ss58_address: Balance.from_tao(999_998),
     }
     assert balances[alice_wallet.coldkey.ss58_address] > alice_balance
+
+
+def test_safe_staking_scenarios(subtensor, alice_wallet, bob_wallet):
+    """
+    Tests safe staking scenarios with different parameters.
+
+    For both staking and unstaking:
+    1. Fails with strict threshold (0.5%) and no partial staking
+    2. Succeeds with strict threshold (0.5%) and partial staking allowed
+    3. Succeeds with lenient threshold (10% and 30%) and no partial staking
+    """
+    netuid = 2
+    # Register root as Alice - the subnet owner and validator
+    assert subtensor.register_subnet(alice_wallet)
+
+    # Verify subnet created successfully
+    assert subtensor.subnet_exists(netuid), "Subnet wasn't created successfully"
+
+    subtensor.burned_register(
+        alice_wallet,
+        netuid=netuid,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+    )
+    subtensor.burned_register(
+        bob_wallet,
+        netuid=netuid,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+    )
+
+    initial_stake = subtensor.get_stake(
+        alice_wallet.coldkey.ss58_address,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+    )
+    assert initial_stake == Balance(0)
+
+    # Test Staking Scenarios
+    stake_amount = Balance.from_tao(100)
+
+    # 1. Strict params - should fail
+    success = subtensor.add_stake(
+        alice_wallet,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+        amount=stake_amount,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        safe_staking=True,
+        rate_threshold=0.005,  # 0.5%
+        allow_partial_stake=False,
+    )
+    assert success is False
+
+    current_stake = subtensor.get_stake(
+        alice_wallet.coldkey.ss58_address,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+    )
+    assert current_stake == Balance(0), "Stake should not change after failed attempt"
+
+    # 2. Partial allowed - should succeed partially
+    success = subtensor.add_stake(
+        alice_wallet,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+        amount=stake_amount,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        safe_staking=True,
+        rate_threshold=0.005,  # 0.5%
+        allow_partial_stake=True,
+    )
+    assert success is True
+
+    partial_stake = subtensor.get_stake(
+        alice_wallet.coldkey.ss58_address,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+    )
+    assert partial_stake > Balance(0), "Partial stake should be added"
+    assert (
+        partial_stake < stake_amount
+    ), "Partial stake should be less than requested amount"
+
+    # 3. Higher threshold - should succeed fully
+    amount = Balance.from_tao(100)
+    success = subtensor.add_stake(
+        alice_wallet,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+        amount=amount,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        safe_staking=True,
+        rate_threshold=0.1,  # 10%
+        allow_partial_stake=False,
+    )
+    assert success is True
+
+    full_stake = subtensor.get_stake(
+        alice_wallet.coldkey.ss58_address,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+    )
+    assert full_stake >= stake_amount, "Full stake amount should be added"
+
+    # Test Unstaking Scenarios
+    # 1. Strict params - should fail
+    success = subtensor.unstake(
+        alice_wallet,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+        amount=stake_amount,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        safe_staking=True,
+        rate_threshold=0.005,  # 0.5%
+        allow_partial_stake=False,
+    )
+    assert success is False
+
+    current_stake = subtensor.get_stake(
+        alice_wallet.coldkey.ss58_address,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+    )
+    assert (
+        current_stake == full_stake
+    ), "Stake should not change after failed unstake attempt"
+
+    # 2. Partial allowed - should succeed partially
+    success = subtensor.unstake(
+        alice_wallet,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+        amount=stake_amount,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        safe_staking=True,
+        rate_threshold=0.005,  # 0.5%
+        allow_partial_stake=True,
+    )
+    assert success is True
+
+    partial_unstake = subtensor.get_stake(
+        alice_wallet.coldkey.ss58_address,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+    )
+    assert partial_unstake > Balance(0), "Some stake should remain"
+
+    # 3. Higher threshold - should succeed fully
+    success = subtensor.unstake(
+        alice_wallet,
+        bob_wallet.hotkey.ss58_address,
+        netuid=netuid,
+        amount=partial_unstake,
+        wait_for_inclusion=True,
+        wait_for_finalization=True,
+        safe_staking=True,
+        rate_threshold=0.3,  # 30%
+        allow_partial_stake=False,
+    )
+    assert success is True