Prevent unstake call from allowing zero amount

pallets/subtensor/src/staking.rs

@@ -189,6 +189,9 @@ impl<T: Config> Pallet<T> {
         // --- 3. Ensure that the hotkey allows delegation or that the hotkey is owned by the calling coldkey.
         ensure!( Self::hotkey_is_delegate( &hotkey ) || Self::coldkey_owns_hotkey( &coldkey, &hotkey ), Error::<T>::NonAssociatedColdKey );
 
+        // --- Ensure that the stake amount to be removed is above zero.
+        ensure!( stake_to_be_removed > 0, Error::<T>::NotEnoughStaketoWithdraw );
+
         // --- 4. Ensure that the hotkey has enough stake to withdraw.
         ensure!( Self::has_enough_stake( &coldkey, &hotkey, stake_to_be_removed ), Error::<T>::NotEnoughStaketoWithdraw );
 

pallets/subtensor/tests/staking.rs

@@ -1,4 +1,4 @@
-use frame_support::{assert_ok, traits::Currency};
+use frame_support::{assert_ok, assert_noop, traits::Currency};
 use frame_system::{Config};
 mod mock;
 use mock::*;
@@ -334,6 +334,38 @@ fn test_remove_stake_ok_no_emission() {
 	});
 }
 
+#[test]
+fn test_remove_stake_amount_zero() {
+	new_test_ext().execute_with(|| {
+		let coldkey_account_id = U256::from(4343);
+		let hotkey_account_id = U256::from(4968585);
+		let amount = 10000;
+        let netuid: u16 = 1;
+		let tempo: u16 = 13;
+		let start_nonce: u64 = 0;
+
+		//add network
+		add_network(netuid, tempo, 0);
+
+		// Let's spin up a neuron
+		register_ok_neuron( netuid, hotkey_account_id, coldkey_account_id, start_nonce);
+
+		// Some basic assertions
+		assert_eq!(SubtensorModule::get_total_stake(), 0);
+		assert_eq!(SubtensorModule::get_total_stake_for_hotkey(&hotkey_account_id), 0);
+		assert_eq!(SubtensorModule::get_coldkey_balance(&coldkey_account_id), 0);
+
+		// Give the neuron some stake to remove
+		SubtensorModule::increase_stake_on_hotkey_account(&hotkey_account_id, amount);
+
+		// Do the magic
+		assert_noop!(
+            SubtensorModule::remove_stake(<<Test as Config>::RuntimeOrigin>::signed(coldkey_account_id), hotkey_account_id, 0),
+            Error::<Test>::NotEnoughStaketoWithdraw
+        );
+	});
+}
+
 #[test]
 fn test_remove_stake_err_signature() {
 	new_test_ext().execute_with(|| {