Fix hardcoded leasing weights

pallets/subtensor/src/subnets/leasing.rs

@@ -16,6 +16,7 @@
 //! ownership will be transferred to the beneficiary.
 
 use super::*;
+use crate::weights::WeightInfo;
 use frame_support::{
     dispatch::RawOrigin,
     traits::{Defensive, fungible::*},
@@ -163,12 +164,10 @@ impl<T: Config> Pallet<T> {
 
         if crowdloan.contributors_count < T::MaxContributors::get() {
             // We have less contributors than the max allowed, so we need to refund the difference
-            Ok(
-                Some(SubnetLeasingWeightInfo::<T>::do_register_leased_network(
-                    crowdloan.contributors_count,
-                ))
-                .into(),
-            )
+            Ok(Some(<T as Config>::WeightInfo::register_leased_network(
+                crowdloan.contributors_count,
+            ))
+            .into())
         } else {
             // We have the max number of contributors, so we don't need to refund anything
             Ok(().into())
@@ -224,10 +223,12 @@ impl<T: Config> Pallet<T> {
             netuid: lease.netuid,
         });
 
-        if clear_result.unique < T::MaxContributors::get() {
+        // Lease shares exclude the beneficiary, while the benchmark's `k` includes them.
+        let contributors_count = clear_result.unique.saturating_add(1);
+        if contributors_count < T::MaxContributors::get() {
             // We have cleared less than the max number of shareholders, so we need to refund the difference
-            Ok(Some(SubnetLeasingWeightInfo::<T>::do_terminate_lease(
-                clear_result.unique,
+            Ok(Some(<T as Config>::WeightInfo::terminate_lease(
+                contributors_count,
             ))
             .into())
         } else {
@@ -378,27 +379,3 @@ impl<T: Config> Pallet<T> {
         Ok((crowdloan_id, crowdloan))
     }
 }
-
-/// Weight functions needed for subnet leasing.
-pub struct SubnetLeasingWeightInfo<T>(PhantomData<T>);
-impl<T: frame_system::Config> SubnetLeasingWeightInfo<T> {
-    pub fn do_register_leased_network(k: u32) -> Weight {
-        Weight::from_parts(301_560_714, 10079)
-            .saturating_add(Weight::from_parts(26_884_006, 0).saturating_mul(k.into()))
-            .saturating_add(T::DbWeight::get().reads(41_u64))
-            .saturating_add(T::DbWeight::get().reads(2_u64.saturating_mul(k.into())))
-            .saturating_add(T::DbWeight::get().writes(55_u64))
-            .saturating_add(T::DbWeight::get().writes(2_u64.saturating_mul(k.into())))
-            .saturating_add(Weight::from_parts(0, 2579).saturating_mul(k.into()))
-    }
-
-    pub fn do_terminate_lease(k: u32) -> Weight {
-        Weight::from_parts(56_635_122, 6148)
-            .saturating_add(Weight::from_parts(912_993, 0).saturating_mul(k.into()))
-            .saturating_add(T::DbWeight::get().reads(4_u64))
-            .saturating_add(T::DbWeight::get().reads((1_u64).saturating_mul(k.into())))
-            .saturating_add(T::DbWeight::get().writes(6_u64))
-            .saturating_add(T::DbWeight::get().writes((1_u64).saturating_mul(k.into())))
-            .saturating_add(Weight::from_parts(0, 2529).saturating_mul(k.into()))
-    }
-}

pallets/subtensor/src/tests/leasing.rs

@@ -28,11 +28,22 @@ fn test_register_leased_network_works() {
         // Register the leased network
         let end_block = 500;
         let emissions_share = Percent::from_percent(30);
-        assert_ok!(SubtensorModule::register_leased_network(
+        let post_info = SubtensorModule::register_leased_network(
             RuntimeOrigin::signed(beneficiary),
             emissions_share,
             Some(end_block),
-        ));
+        )
+        .expect("leased network registration should succeed");
+
+        let contributors_count = 1 + contributions.len() as u32;
+        assert_eq!(
+            post_info.actual_weight,
+            Some(
+                <<Test as crate::Config>::WeightInfo as crate::weights::WeightInfo>::register_leased_network(
+                    contributors_count,
+                ),
+            )
+        );
 
         // Ensure the lease was created
         let lease_id = 0;
@@ -239,25 +250,33 @@ fn test_register_lease_network_fails_if_end_block_is_in_the_past() {
 
 #[test]
 fn test_terminate_lease_works() {
-    new_test_ext(1).execute_with(|| {
+    let mut ext = new_test_ext(1);
+    let beneficiary = U256::from(1);
+    let contributions = vec![(U256::from(2), 990_000_000_000)]; // 990 TAO
+    let end_block = 500;
+
+    let (lease_id, lease) = ext.execute_with(|| {
         // Setup a crowdloan
         let crowdloan_id = 0;
-        let beneficiary = U256::from(1);
         let deposit = 10_000_000_000; // 10 TAO
         let cap = 1_000_000_000_000; // 1000 TAO
-        let contributions = vec![(U256::from(2), 990_000_000_000)]; // 990 TAO
         setup_crowdloan(crowdloan_id, deposit, cap, beneficiary, &contributions);
 
         // Setup a leased network
-        let end_block = 500;
         let tao_to_stake = 100_000_000_000; // 100 TAO
         let emissions_share = Percent::from_percent(30);
-        let (lease_id, lease) = setup_leased_network(
+        setup_leased_network(
             beneficiary,
             emissions_share,
             Some(end_block),
             Some(tao_to_stake),
-        );
+        )
+    });
+
+    // Commit the lease setup so clear_prefix reports the same backend removals as on-chain.
+    ext.commit_all().expect("lease setup should commit");
+
+    ext.execute_with(|| {
 
         // Run to the end of the lease
         run_to_block(end_block);
@@ -266,12 +285,28 @@ fn test_terminate_lease_works() {
         let hotkey = U256::from(3);
         let _ = SubtensorModule::create_account_if_non_existent(&beneficiary, &hotkey);
 
+        assert_eq!(
+            SubnetLeaseShares::<Test>::iter_prefix(lease_id).count(),
+            contributions.len()
+        );
+
         // Terminate the lease
-        assert_ok!(SubtensorModule::terminate_lease(
+        let post_info = SubtensorModule::terminate_lease(
             RuntimeOrigin::signed(beneficiary),
             lease_id,
             hotkey,
-        ));
+        )
+        .expect("lease termination should succeed");
+
+        let contributors_count = 1 + contributions.len() as u32;
+        assert_eq!(
+            post_info.actual_weight,
+            Some(
+                <<Test as crate::Config>::WeightInfo as crate::weights::WeightInfo>::terminate_lease(
+                    contributors_count,
+                ),
+            )
+        );
 
         // Ensure the beneficiary is now the owner of the subnet
         assert_eq!(SubnetOwner::<Test>::get(lease.netuid), beneficiary);