Fix balance updates for add/remove liquidity

pallets/subtensor/src/lib.rs

@@ -2448,7 +2448,7 @@ impl<T: Config + pallet_balances::Config<Balance = u64>>
         pallet_balances::Pallet::<T>::free_balance(account_id)
     }
 
-    fn alpha_balance(netuid: u16, account_id: &T::AccountId) -> u64 {
-        TotalHotkeyAlpha::<T>::get(account_id, netuid)
+    fn alpha_balance(netuid: u16, coldkey: &T::AccountId, hotkey: &T::AccountId) -> u64 {
+        Self::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, netuid)
     }
 }

pallets/subtensor/src/tests/mock.rs

@@ -430,8 +430,8 @@ impl LiquidityDataProvider<AccountId> for SubtensorModule {
         Balances::free_balance(account_id)
     }
 
-    fn alpha_balance(netuid: u16, account_id: &AccountId) -> u64 {
-        TotalHotkeyAlpha::<Test>::get(account_id, netuid)
+    fn alpha_balance(netuid: u16, coldkey: &AccountId, hotkey: &AccountId) -> u64 {
+        SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet(hotkey, coldkey, netuid)
     }
 }
 

pallets/swap-interface/src/lib.rs

@@ -36,5 +36,5 @@ pub trait LiquidityDataProvider<AccountId> {
     fn tao_reserve(netuid: u16) -> u64;
     fn alpha_reserve(netuid: u16) -> u64;
     fn tao_balance(account_id: &AccountId) -> u64;
-    fn alpha_balance(netuid: u16, account_id: &AccountId) -> u64;
+    fn alpha_balance(netuid: u16, coldkey_account_id: &AccountId, hotkey_account_id: &AccountId) -> u64;
 }

pallets/swap/Cargo.toml

@@ -21,6 +21,7 @@ sp-std = { workspace = true }
 substrate-fixed = { workspace = true }
 
 subtensor-swap-interface = { workspace = true }
+pallet-subtensor = { version = "4.0.0-dev", default-features = false, path = "../subtensor" }
 
 [lints]
 workspace = true
@@ -33,6 +34,7 @@ std = [
     "frame-benchmarking/std",
     "frame-support/std",
     "frame-system/std",
+    "pallet-subtensor/std",
     "subtensor-swap-interface/std",
     "safe-math/std",
     "scale-info/std",

pallets/swap/src/mock.rs

@@ -20,7 +20,8 @@ construct_runtime!(
 
 pub type Block = frame_system::mocking::MockBlock<Test>;
 pub type AccountId = u32;
-pub const OK_ACCOUNT_ID: AccountId = 1;
+pub const OK_COLDKEY_ACCOUNT_ID: AccountId = 1;
+pub const OK_HOTKEY_ACCOUNT_ID: AccountId = 1000;
 
 parameter_types! {
     pub const BlockHashCount: u64 = 250;
@@ -70,24 +71,30 @@ parameter_types! {
 pub struct MockLiquidityProvider;
 
 impl LiquidityDataProvider<AccountId> for MockLiquidityProvider {
-    fn tao_reserve(_: u16) -> u64 {
-        1_000_000_000_000
+    fn tao_reserve(netuid: u16) -> u64 {
+        match netuid {
+            123 => 1_000,
+            _ => 1_000_000_000_000
+        }
     }
 
-    fn alpha_reserve(_: u16) -> u64 {
-        4_000_000_000_000
+    fn alpha_reserve(netuid: u16) -> u64 {
+        match netuid {
+            123 => 1,
+            _ => 4_000_000_000_000
+        }
     }
 
     fn tao_balance(account_id: &AccountId) -> u64 {
-        if *account_id == OK_ACCOUNT_ID {
+        if *account_id == OK_COLDKEY_ACCOUNT_ID {
             100_000_000_000_000
         } else {
             1_000_000_000
         }
     }
 
-    fn alpha_balance(_: u16, account_id: &AccountId) -> u64 {
-        if *account_id == OK_ACCOUNT_ID {
+    fn alpha_balance(_: u16, coldkey_account_id: &AccountId, hotkey_account_id: &AccountId) -> u64 {
+        if (*coldkey_account_id == OK_COLDKEY_ACCOUNT_ID) && (*hotkey_account_id == OK_HOTKEY_ACCOUNT_ID) {
             100_000_000_000_000
         } else {
             1_000_000_000