Add eco-tests for Mentat

chain-extensions/src/tests.rs

@@ -1297,15 +1297,26 @@ fn add_stake_recycle_rollback_on_recycle_failure() {
 
         let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
-        // Set up very low reserves so recycle will fail with InsufficientLiquidity
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::insert_lock_state(
+            &coldkey,
+            netuid,
+            &hotkey,
+            pallet_subtensor::staking::lock::LockState {
+                locked_mass: AlphaBalance::from(u64::MAX / 4),
+                conviction: U64F64::saturating_from_num(0),
+                last_update: pallet_subtensor::Pallet::<mock::Test>::get_current_block_as_u64(),
+            },
+        );
+
+        // Leave enough input-side liquidity for add_stake to pass the 1000x swap input cap.
+        // The lock above makes the recycle leg fail, exercising atomic rollback.
         mock::setup_reserves(
             netuid,
-            TaoBalance::from(1_000_u64),
+            TaoBalance::from(tao_amount_raw / 1000 + 1),
             AlphaBalance::from(1_000_u64),
         );
 
-        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
-
         add_balance_to_coldkey_account(
             &coldkey,
             TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),
@@ -1368,15 +1379,26 @@ fn add_stake_burn_rollback_on_burn_failure() {
 
         let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
 
-        // Set up very low reserves so burn will fail with InsufficientLiquidity
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::insert_lock_state(
+            &coldkey,
+            netuid,
+            &hotkey,
+            pallet_subtensor::staking::lock::LockState {
+                locked_mass: AlphaBalance::from(u64::MAX / 4),
+                conviction: U64F64::saturating_from_num(0),
+                last_update: pallet_subtensor::Pallet::<mock::Test>::get_current_block_as_u64(),
+            },
+        );
+
+        // Leave enough input-side liquidity for add_stake to pass the 1000x swap input cap.
+        // The lock above makes the burn leg fail, exercising atomic rollback.
         mock::setup_reserves(
             netuid,
-            TaoBalance::from(1_000_u64),
+            TaoBalance::from(tao_amount_raw / 1000 + 1),
             AlphaBalance::from(1_000_u64),
         );
 
-        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
-
         add_balance_to_coldkey_account(
             &coldkey,
             TaoBalance::from(tao_amount_raw.saturating_add(1_000_000_000)),

eco-tests/src/lib.rs

@@ -7,3 +7,6 @@ mod tests;
 
 #[cfg(test)]
 mod tests_taocom_indexer;
+
+#[cfg(test)]
+mod tests_mentat_indexer;

eco-tests/src/mock.rs

@@ -612,11 +612,19 @@ pub fn add_balance_to_coldkey_account(coldkey: &U256, tao: TaoBalance) {
 mod api_mocks {
     use codec::Compact;
     use pallet_subtensor::rpc_info::delegate_info::DelegateInfo;
+    use pallet_subtensor::rpc_info::dynamic_info::DynamicInfo;
+    use pallet_subtensor::rpc_info::metagraph::{Metagraph, SelectiveMetagraph};
+    use pallet_subtensor::rpc_info::show_subnet::SubnetState;
     use pallet_subtensor::rpc_info::stake_info::StakeInfo;
+    use pallet_subtensor::rpc_info::subnet_info::{
+        SubnetHyperparams, SubnetHyperparamsV2, SubnetHyperparamsV3, SubnetInfo, SubnetInfov2,
+    };
     use pallet_subtensor_swap_runtime_api::{SimSwapResult, SubnetPrice, SwapRuntimeApi};
     use sp_runtime::AccountId32;
-    use subtensor_custom_rpc_runtime_api::{DelegateInfoRuntimeApi, StakeInfoRuntimeApi};
-    use subtensor_runtime_common::{AlphaBalance, NetUid, TaoBalance};
+    use subtensor_custom_rpc_runtime_api::{
+        DelegateInfoRuntimeApi, StakeInfoRuntimeApi, SubnetInfoRuntimeApi,
+    };
+    use subtensor_runtime_common::{AlphaBalance, MechId, NetUid, TaoBalance};
 
     use super::Block;
 
@@ -660,6 +668,31 @@ mod api_mocks {
             }
         }
 
+        impl SubnetInfoRuntimeApi<Block> for MockApi {
+            fn get_subnet_info(_netuid: NetUid) -> Option<SubnetInfo<AccountId32>> { None }
+            fn get_subnets_info() -> Vec<Option<SubnetInfo<AccountId32>>> { Vec::new() }
+            fn get_subnet_info_v2(_netuid: NetUid) -> Option<SubnetInfov2<AccountId32>> { None }
+            fn get_subnets_info_v2() -> Vec<Option<SubnetInfov2<AccountId32>>> { Vec::new() }
+            #[allow(deprecated)]
+            fn get_subnet_hyperparams(_netuid: NetUid) -> Option<SubnetHyperparams> { None }
+            #[allow(deprecated)]
+            fn get_subnet_hyperparams_v2(_netuid: NetUid) -> Option<SubnetHyperparamsV2> { None }
+            fn get_subnet_hyperparams_v3(_netuid: NetUid) -> Option<SubnetHyperparamsV3> { None }
+            fn get_all_dynamic_info() -> Vec<Option<DynamicInfo<AccountId32>>> { Vec::new() }
+            fn get_all_metagraphs() -> Vec<Option<Metagraph<AccountId32>>> { Vec::new() }
+            fn get_metagraph(_netuid: NetUid) -> Option<Metagraph<AccountId32>> { None }
+            fn get_all_mechagraphs() -> Vec<Option<Metagraph<AccountId32>>> { Vec::new() }
+            fn get_mechagraph(_netuid: NetUid, _mecid: MechId) -> Option<Metagraph<AccountId32>> { None }
+            fn get_dynamic_info(_netuid: NetUid) -> Option<DynamicInfo<AccountId32>> { None }
+            fn get_subnet_state(_netuid: NetUid) -> Option<SubnetState<AccountId32>> { None }
+            fn get_selective_metagraph(_netuid: NetUid, _metagraph_indexes: Vec<u16>) -> Option<SelectiveMetagraph<AccountId32>> { None }
+            fn get_coldkey_auto_stake_hotkey(_coldkey: AccountId32, _netuid: NetUid) -> Option<AccountId32> { None }
+            fn get_selective_mechagraph(_netuid: NetUid, _subid: MechId, _metagraph_indexes: Vec<u16>) -> Option<SelectiveMetagraph<AccountId32>> { None }
+            fn get_subnet_to_prune() -> Option<NetUid> { None }
+            fn get_subnet_account_id(_netuid: NetUid) -> Option<AccountId32> { None }
+            fn get_next_epoch_start_block(_netuid: NetUid) -> Option<u64> { None }
+        }
+
         impl SwapRuntimeApi<Block> for MockApi {
             fn current_alpha_price(_netuid: NetUid) -> u64 { 0 }
             fn current_alpha_price_all() -> Vec<SubnetPrice> { Vec::new() }