diff --git a/.github/workflows/check-bittensor-e2e-tests.yml.yml b/.github/workflows/check-bittensor-e2e-tests.yml.yml index 540238d6d5..c133efab22 100644 --- a/.github/workflows/check-bittensor-e2e-tests.yml.yml +++ b/.github/workflows/check-bittensor-e2e-tests.yml.yml @@ -130,17 +130,27 @@ jobs: - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 + - name: Move Docker data-root to /mnt/data + run: | + sudo systemctl stop docker + sudo mkdir -p /mnt/data/docker + sudo chown -R runner:runner /mnt/data + sudo chmod -R 777 /mnt/data + echo '{"data-root": "/mnt/data/docker"}' | sudo tee /etc/docker/daemon.json + sudo systemctl start docker + docker info | grep "Docker Root Dir" + - name: Build Docker Image run: docker build -f Dockerfile-localnet -t localnet . - name: Save Docker Image as Tar - run: docker save -o subtensor-localnet.tar localnet + run: docker save -o /mnt/data/subtensor-localnet.tar localnet - name: Upload Docker Image as Artifact uses: actions/upload-artifact@v4 with: name: subtensor-localnet - path: subtensor-localnet.tar + path: /mnt/data/subtensor-localnet.tar # main btcli job run-btcli-e2e-tests: @@ -181,7 +191,7 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v5 with: - enable-cache: 'false' + enable-cache: "false" - name: Create Python virtual environment working-directory: ${{ github.workspace }} @@ -275,7 +285,7 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v5 with: - enable-cache: 'false' + enable-cache: "false" - name: Create Python virtual environment working-directory: ${{ github.workspace }} diff --git a/evm-tests/test/staking.precompile.reward.test.ts b/evm-tests/test/staking.precompile.reward.test.ts index 1b47239e1d..d04620c91b 100644 --- a/evm-tests/test/staking.precompile.reward.test.ts +++ b/evm-tests/test/staking.precompile.reward.test.ts @@ -74,7 +74,7 @@ describe("Test neuron precompile reward", () => { let index = 0; while (index < 60) { - const pending = await api.query.SubtensorModule.ValidatorEmission.getValue(netuid); + const pending = await api.query.SubtensorModule.PendingValidatorEmission.getValue(netuid); if (pending > 0) { console.log("pending amount is ", pending); break; diff --git a/pallets/subtensor/src/macros/events.rs b/pallets/subtensor/src/macros/events.rs index c2931024ee..d015205d4d 100644 --- a/pallets/subtensor/src/macros/events.rs +++ b/pallets/subtensor/src/macros/events.rs @@ -467,5 +467,15 @@ mod events { /// Claim type root_claim_type: RootClaimTypeEnum, }, + + /// Subnet lease dividends have been distributed. + SubnetLeaseDividendsDistributed { + /// The lease ID + lease_id: LeaseId, + /// The contributor + contributor: T::AccountId, + /// The amount of alpha distributed + alpha: AlphaCurrency, + }, } } diff --git a/pallets/subtensor/src/subnets/leasing.rs b/pallets/subtensor/src/subnets/leasing.rs index cf263a1335..a202a22a45 100644 --- a/pallets/subtensor/src/subnets/leasing.rs +++ b/pallets/subtensor/src/subnets/leasing.rs @@ -24,8 +24,7 @@ use frame_system::pallet_prelude::*; use sp_core::blake2_256; use sp_runtime::{Percent, traits::TrailingZeroInput}; use substrate_fixed::types::U64F64; -use subtensor_runtime_common::{AlphaCurrency, NetUid, TaoCurrency}; -use subtensor_swap_interface::SwapHandler; +use subtensor_runtime_common::{AlphaCurrency, NetUid}; pub type LeaseId = u32; @@ -130,6 +129,9 @@ impl Pallet { ); SubnetUidToLeaseId::::insert(netuid, lease_id); + // The lease take should be 0% to allow all contributors to receive dividends entirely. + Self::delegate_hotkey(&lease_hotkey, 0); + // Get all the contributions to the crowdloan except for the beneficiary // because its share will be computed as the dividends are distributed let contributions = pallet_crowdloan::Contributions::::iter_prefix(crowdloan_id) @@ -249,9 +251,8 @@ impl Pallet { /// Hook used when the subnet owner's cut is distributed to split the amount into dividends /// for the contributors and the beneficiary in shares relative to their initial contributions. - /// - /// It will ensure the subnet has enough alpha in its liquidity pool before swapping it to tao to be distributed, - /// and if not enough liquidity is available, it will accumulate the dividends for later distribution. + /// It accumulates dividends to be distributed later when the interval for distribution is reached. + /// Distribution is made in alpha and stake to the contributor coldkey and lease hotkey. pub fn distribute_leased_network_dividends(lease_id: LeaseId, owner_cut_alpha: AlphaCurrency) { // Ensure the lease exists let Some(lease) = SubnetLeases::::get(lease_id) else { @@ -290,55 +291,60 @@ impl Pallet { return; } - // Ensure there is enough liquidity to unstake the contributors cut - if let Err(err) = Self::validate_remove_stake( - &lease.coldkey, - &lease.hotkey, - lease.netuid, - total_contributors_cut_alpha, - total_contributors_cut_alpha, - false, - ) { - log::debug!("Couldn't distributing dividends for lease {lease_id}: {err:?}"); - AccumulatedLeaseDividends::::set(lease_id, total_contributors_cut_alpha); - return; - } - - // Unstake the contributors cut from the subnet as tao to the lease coldkey - let tao_unstaked = match Self::unstake_from_subnet( - &lease.hotkey, - &lease.coldkey, - lease.netuid, - total_contributors_cut_alpha, - T::SwapInterface::min_price(), - false, - ) { - Ok(tao_unstaked) => tao_unstaked, - Err(err) => { - log::debug!("Couldn't distributing dividends for lease {lease_id}: {err:?}"); - AccumulatedLeaseDividends::::set(lease_id, total_contributors_cut_alpha); - return; + // We use a storage layer to ensure the distribution is atomic. + if let Err(err) = frame_support::storage::with_storage_layer(|| { + let mut alpha_distributed = AlphaCurrency::ZERO; + + // Distribute the contributors cut to the contributors and accumulate the alpha + // distributed so far to obtain how much alpha is left to distribute to the beneficiary + for (contributor, share) in SubnetLeaseShares::::iter_prefix(lease_id) { + let alpha_for_contributor = share + .saturating_mul(U64F64::from(total_contributors_cut_alpha.to_u64())) + .ceil() + .saturating_to_num::(); + + Self::transfer_stake_within_subnet( + &lease.coldkey, + &lease.hotkey, + &contributor, + &lease.hotkey, + lease.netuid, + alpha_for_contributor.into(), + )?; + alpha_distributed = alpha_distributed.saturating_add(alpha_for_contributor.into()); + + Self::deposit_event(Event::SubnetLeaseDividendsDistributed { + lease_id, + contributor, + alpha: alpha_for_contributor.into(), + }); } - }; - // Distribute the contributors cut to the contributors and accumulate the tao - // distributed so far to obtain how much tao is left to distribute to the beneficiary - let mut tao_distributed = TaoCurrency::ZERO; - for (contributor, share) in SubnetLeaseShares::::iter_prefix(lease_id) { - let tao_for_contributor = share - .saturating_mul(U64F64::from(tao_unstaked.to_u64())) - .floor() - .saturating_to_num::(); - Self::add_balance_to_coldkey_account(&contributor, tao_for_contributor); - tao_distributed = tao_distributed.saturating_add(tao_for_contributor.into()); - } + // Distribute the leftover alpha to the beneficiary + let beneficiary_cut_alpha = + total_contributors_cut_alpha.saturating_sub(alpha_distributed); + Self::transfer_stake_within_subnet( + &lease.coldkey, + &lease.hotkey, + &lease.beneficiary, + &lease.hotkey, + lease.netuid, + beneficiary_cut_alpha.into(), + )?; + Self::deposit_event(Event::SubnetLeaseDividendsDistributed { + lease_id, + contributor: lease.beneficiary.clone(), + alpha: beneficiary_cut_alpha.into(), + }); - // Distribute the leftover tao to the beneficiary - let beneficiary_cut_tao = tao_unstaked.saturating_sub(tao_distributed); - Self::add_balance_to_coldkey_account(&lease.beneficiary, beneficiary_cut_tao.into()); + // Reset the accumulated dividends + AccumulatedLeaseDividends::::insert(lease_id, AlphaCurrency::ZERO); - // Reset the accumulated dividends - AccumulatedLeaseDividends::::insert(lease_id, AlphaCurrency::ZERO); + Ok::<(), DispatchError>(()) + }) { + log::debug!("Couldn't distributing dividends for lease {lease_id}: {err:?}"); + AccumulatedLeaseDividends::::set(lease_id, total_contributors_cut_alpha); + }; } fn lease_coldkey(lease_id: LeaseId) -> Result { diff --git a/pallets/subtensor/src/tests/leasing.rs b/pallets/subtensor/src/tests/leasing.rs index 4ffc18230a..9f8bf4d6bf 100644 --- a/pallets/subtensor/src/tests/leasing.rs +++ b/pallets/subtensor/src/tests/leasing.rs @@ -65,6 +65,9 @@ fn test_register_leased_network_works() { contributor2_share ); + // Ensure the lease hotkey has 0 take from staking + assert_eq!(SubtensorModule::get_hotkey_take(&lease.hotkey), 0); + // Ensure each contributor and beneficiary has been refunded their share of the leftover cap let leftover_cap = cap.saturating_sub(lease.cost); @@ -502,54 +505,111 @@ fn test_distribute_lease_network_dividends_multiple_contributors_works() { // Setup the correct block to distribute dividends run_to_block(::LeaseDividendsDistributionInterval::get() as u64); - // Get the initial subnet tao after stake and ensure all contributor - // balances are in initial state - let subnet_tao_before = SubnetTAO::::get(lease.netuid); - let contributor1_balance_before = SubtensorModule::get_coldkey_balance(&contributions[0].0); - let contributor2_balance_before = SubtensorModule::get_coldkey_balance(&contributions[1].0); - let beneficiary_balance_before = SubtensorModule::get_coldkey_balance(&beneficiary); + // Get the initial alpha for the contributors and beneficiary and ensure they are zero + let contributor1_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid, + ); + assert_eq!(contributor1_alpha_before, AlphaCurrency::ZERO); + let contributor2_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid, + ); + assert_eq!(contributor2_alpha_before, AlphaCurrency::ZERO); + let beneficiary_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ); + assert_eq!(beneficiary_alpha_before, AlphaCurrency::ZERO); // Setup some previously accumulated dividends - let accumulated_dividends = AlphaCurrency::from(5_000_000); + let accumulated_dividends = AlphaCurrency::from(10_000_000_000); AccumulatedLeaseDividends::::insert(lease_id, accumulated_dividends); // Distribute the dividends - let owner_cut_alpha = AlphaCurrency::from(5_000_000); + let owner_cut_alpha = AlphaCurrency::from(5_000_000_000); SubtensorModule::distribute_leased_network_dividends(lease_id, owner_cut_alpha); // Ensure the dividends were distributed correctly relative to their shares - let distributed_tao = subnet_tao_before - SubnetTAO::::get(lease.netuid); - let contributor1_balance_delta = SubtensorModule::get_coldkey_balance(&contributions[0].0) - .saturating_sub(contributor1_balance_before); - let contributor2_balance_delta = SubtensorModule::get_coldkey_balance(&contributions[1].0) - .saturating_sub(contributor2_balance_before); - let beneficiary_balance_delta = SubtensorModule::get_coldkey_balance(&beneficiary) - .saturating_sub(beneficiary_balance_before); - + let distributed_alpha = + accumulated_dividends + emissions_share.mul_ceil(owner_cut_alpha.to_u64()).into(); + assert_ne!(distributed_alpha, AlphaCurrency::ZERO); + + let contributor1_alpha_delta = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid, + ) + .saturating_sub(contributor1_alpha_before); + assert_ne!(contributor1_alpha_delta, AlphaCurrency::ZERO); + + let contributor2_alpha_delta = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid, + ) + .saturating_sub(contributor2_alpha_before); + assert_ne!(contributor2_alpha_delta, AlphaCurrency::ZERO); + + let beneficiary_alpha_delta = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ) + .saturating_sub(beneficiary_alpha_before); + assert_ne!(beneficiary_alpha_delta, AlphaCurrency::ZERO); + + // What has been distributed should be equal to the sum of all contributors received alpha assert_eq!( - distributed_tao, - (beneficiary_balance_delta + contributor1_balance_delta + contributor2_balance_delta) - .into() + distributed_alpha, + (beneficiary_alpha_delta + contributor1_alpha_delta + contributor2_alpha_delta).into() ); - let expected_contributor1_balance = + let expected_contributor1_alpha = SubnetLeaseShares::::get(lease_id, contributions[0].0) - .saturating_mul(U64F64::from(distributed_tao.to_u64())) - .floor() + .saturating_mul(U64F64::from(distributed_alpha.to_u64())) + .ceil() .to_num::(); - assert_eq!(contributor1_balance_delta, expected_contributor1_balance); + assert_eq!(contributor1_alpha_delta, expected_contributor1_alpha.into()); + assert_eq!( + System::events()[2].event, + RuntimeEvent::SubtensorModule(Event::SubnetLeaseDividendsDistributed { + lease_id, + contributor: contributions[0].0.into(), + alpha: expected_contributor1_alpha.into(), + },) + ); - let expected_contributor2_balance = + let expected_contributor2_alpha = SubnetLeaseShares::::get(lease_id, contributions[1].0) - .saturating_mul(U64F64::from(distributed_tao.to_u64())) - .floor() + .saturating_mul(U64F64::from(distributed_alpha.to_u64())) + .ceil() .to_num::(); - assert_eq!(contributor2_balance_delta, expected_contributor2_balance); + assert_eq!(contributor2_alpha_delta, expected_contributor2_alpha.into()); + assert_eq!( + System::events()[5].event, + RuntimeEvent::SubtensorModule(Event::SubnetLeaseDividendsDistributed { + lease_id, + contributor: contributions[1].0.into(), + alpha: expected_contributor2_alpha.into(), + },) + ); // The beneficiary should have received the remaining dividends - let expected_beneficiary_balance = distributed_tao.to_u64() - - (expected_contributor1_balance + expected_contributor2_balance); - assert_eq!(beneficiary_balance_delta, expected_beneficiary_balance); + let expected_beneficiary_alpha = distributed_alpha.to_u64() + - (expected_contributor1_alpha + expected_contributor2_alpha); + assert_eq!(beneficiary_alpha_delta, expected_beneficiary_alpha.into()); + assert_eq!( + System::events()[8].event, + RuntimeEvent::SubtensorModule(Event::SubnetLeaseDividendsDistributed { + lease_id, + contributor: beneficiary.into(), + alpha: expected_beneficiary_alpha.into(), + },) + ); // Ensure nothing was accumulated for later distribution assert_eq!( @@ -584,23 +644,40 @@ fn test_distribute_lease_network_dividends_only_beneficiary_works() { // Setup the correct block to distribute dividends run_to_block(::LeaseDividendsDistributionInterval::get() as u64); - // Get the initial subnet tao after stake and beneficiary balance - let subnet_tao_before = SubnetTAO::::get(lease.netuid); - let beneficiary_balance_before = SubtensorModule::get_coldkey_balance(&beneficiary); + // Get the initial alpha for the beneficiary and ensure it is zero + let beneficiary_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ); + assert_eq!(beneficiary_alpha_before, AlphaCurrency::ZERO); // Setup some previously accumulated dividends - let accumulated_dividends = AlphaCurrency::from(5_000_000); + let accumulated_dividends = AlphaCurrency::from(10_000_000_000); AccumulatedLeaseDividends::::insert(lease_id, accumulated_dividends); // Distribute the dividends - let owner_cut_alpha = AlphaCurrency::from(5_000_000); + let owner_cut_alpha = AlphaCurrency::from(5_000_000_000); SubtensorModule::distribute_leased_network_dividends(lease_id, owner_cut_alpha); // Ensure the dividends were distributed correctly relative to their shares - let distributed_tao = subnet_tao_before - SubnetTAO::::get(lease.netuid); - let beneficiary_balance_delta = SubtensorModule::get_coldkey_balance(&beneficiary) - .saturating_sub(beneficiary_balance_before); - assert_eq!(distributed_tao, beneficiary_balance_delta.into()); + let distributed_alpha = + accumulated_dividends + emissions_share.mul_ceil(owner_cut_alpha.to_u64()).into(); + assert_ne!(distributed_alpha, AlphaCurrency::ZERO); + let beneficiary_alpha_delta = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ) + .saturating_sub(beneficiary_alpha_before); + assert_eq!(beneficiary_alpha_delta, distributed_alpha.into()); + assert_last_event::(RuntimeEvent::SubtensorModule( + Event::SubnetLeaseDividendsDistributed { + lease_id, + contributor: beneficiary.into(), + alpha: distributed_alpha, + }, + )); // Ensure nothing was accumulated for later distribution assert_eq!( @@ -628,7 +705,7 @@ fn test_distribute_lease_network_dividends_accumulates_if_not_the_correct_block( let end_block = 500; let emissions_share = Percent::from_percent(30); let tao_to_stake = 100_000_000_000; // 100 TAO - let (lease_id, _) = setup_leased_network( + let (lease_id, lease) = setup_leased_network( beneficiary, emissions_share, Some(end_block), @@ -638,31 +715,58 @@ fn test_distribute_lease_network_dividends_accumulates_if_not_the_correct_block( // Setup incorrect block to distribute dividends run_to_block(::LeaseDividendsDistributionInterval::get() as u64 + 1); - // Get the initial subnet tao after stake and ensure all contributor - let contributor1_balance_before = SubtensorModule::get_coldkey_balance(&contributions[0].0); - let contributor2_balance_before = SubtensorModule::get_coldkey_balance(&contributions[1].0); - let beneficiary_balance_before = SubtensorModule::get_coldkey_balance(&beneficiary); + // Get the initial alpha for the contributors and beneficiary and ensure they are zero + let contributor1_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid, + ); + assert_eq!(contributor1_alpha_before, AlphaCurrency::ZERO); + let contributor2_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid, + ); + assert_eq!(contributor2_alpha_before, AlphaCurrency::ZERO); + let beneficiary_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ); + assert_eq!(beneficiary_alpha_before, AlphaCurrency::ZERO); // Setup some previously accumulated dividends - let accumulated_dividends = AlphaCurrency::from(5_000_000); + let accumulated_dividends = AlphaCurrency::from(10_000_000_000); AccumulatedLeaseDividends::::insert(lease_id, accumulated_dividends); // Distribute the dividends - let owner_cut_alpha = AlphaCurrency::from(5_000_000); + let owner_cut_alpha = AlphaCurrency::from(5_000_000_000); SubtensorModule::distribute_leased_network_dividends(lease_id, owner_cut_alpha); // Ensure the dividends were not distributed assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[0].0), - contributor1_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid + ), + contributor1_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[1].0), - contributor2_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid + ), + contributor2_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&beneficiary), - beneficiary_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid + ), + beneficiary_alpha_before ); // Ensure we correctly accumulated the dividends @@ -711,29 +815,58 @@ fn test_distribute_lease_network_dividends_does_nothing_if_lease_has_ended() { // Run to the end of the lease run_to_block(end_block); - let subnet_tao_before = SubnetTAO::::get(lease.netuid); - let contributor1_balance_before = SubtensorModule::get_coldkey_balance(&contributions[0].0); - let contributor2_balance_before = SubtensorModule::get_coldkey_balance(&contributions[1].0); - let beneficiary_balance_before = SubtensorModule::get_coldkey_balance(&beneficiary); + // Get the initial alpha for the contributors and beneficiary and ensure they are zero + let contributor1_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid, + ); + assert_eq!(contributor1_alpha_before, AlphaCurrency::ZERO); + let contributor2_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid, + ); + assert_eq!(contributor2_alpha_before, AlphaCurrency::ZERO); + let beneficiary_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ); + assert_eq!(beneficiary_alpha_before, AlphaCurrency::ZERO); + + // No dividends are present, lease is new let accumulated_dividends_before = AccumulatedLeaseDividends::::get(lease_id); + assert_eq!(accumulated_dividends_before, AlphaCurrency::ZERO); // Try to distribute the dividends - let owner_cut_alpha = AlphaCurrency::from(5_000_000); + let owner_cut_alpha = AlphaCurrency::from(5_000_000_000); SubtensorModule::distribute_leased_network_dividends(lease_id, owner_cut_alpha); // Ensure the dividends were not distributed - assert_eq!(SubnetTAO::::get(lease.netuid), subnet_tao_before); assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[0].0), - contributor1_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid + ), + contributor1_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[1].0), - contributor2_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid + ), + contributor2_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&beneficiary), - beneficiary_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid + ), + beneficiary_alpha_before ); // Ensure nothing was accumulated for later distribution assert_eq!( @@ -768,28 +901,54 @@ fn test_distribute_lease_network_dividends_accumulates_if_amount_is_too_low() { None, // We don't add any liquidity ); - let subnet_tao_before = SubnetTAO::::get(lease.netuid); - let contributor1_balance_before = SubtensorModule::get_coldkey_balance(&contributions[0].0); - let contributor2_balance_before = SubtensorModule::get_coldkey_balance(&contributions[1].0); - let beneficiary_balance_before = SubtensorModule::get_coldkey_balance(&beneficiary); + // Get the initial alpha for the contributors and beneficiary and ensure they are zero + let contributor1_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid, + ); + assert_eq!(contributor1_alpha_before, AlphaCurrency::ZERO); + let contributor2_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid, + ); + assert_eq!(contributor2_alpha_before, AlphaCurrency::ZERO); + let beneficiary_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ); + assert_eq!(beneficiary_alpha_before, AlphaCurrency::ZERO); // Try to distribute the dividends let owner_cut_alpha = AlphaCurrency::from(5_000); SubtensorModule::distribute_leased_network_dividends(lease_id, owner_cut_alpha); // Ensure the dividends were not distributed - assert_eq!(SubnetTAO::::get(lease.netuid), subnet_tao_before); assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[0].0), - contributor1_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid + ), + contributor1_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[1].0), - contributor2_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid + ), + contributor2_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&beneficiary), - beneficiary_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid + ), + beneficiary_alpha_before ); // Ensure the correct amount of alpha was accumulated for later dividends distribution assert_eq!( @@ -824,28 +983,53 @@ fn test_distribute_lease_network_dividends_accumulates_if_insufficient_liquidity None, // We don't add any liquidity ); - let subnet_tao_before = SubnetTAO::::get(lease.netuid); - let contributor1_balance_before = SubtensorModule::get_coldkey_balance(&contributions[0].0); - let contributor2_balance_before = SubtensorModule::get_coldkey_balance(&contributions[1].0); - let beneficiary_balance_before = SubtensorModule::get_coldkey_balance(&beneficiary); + let contributor1_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid, + ); + assert_eq!(contributor1_alpha_before, AlphaCurrency::ZERO); + let contributor2_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid, + ); + assert_eq!(contributor2_alpha_before, AlphaCurrency::ZERO); + let beneficiary_alpha_before = SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid, + ); + assert_eq!(beneficiary_alpha_before, AlphaCurrency::ZERO); // Try to distribute the dividends let owner_cut_alpha = AlphaCurrency::from(5_000_000); SubtensorModule::distribute_leased_network_dividends(lease_id, owner_cut_alpha); // Ensure the dividends were not distributed - assert_eq!(SubnetTAO::::get(lease.netuid), subnet_tao_before); assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[0].0), - contributor1_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[0].0, + lease.netuid + ), + contributor1_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&contributions[1].0), - contributor2_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &contributions[1].0, + lease.netuid + ), + contributor2_alpha_before ); assert_eq!( - SubtensorModule::get_coldkey_balance(&beneficiary), - beneficiary_balance_before + SubtensorModule::get_stake_for_hotkey_and_coldkey_on_subnet( + &lease.hotkey, + &beneficiary, + lease.netuid + ), + beneficiary_alpha_before ); // Ensure the correct amount of alpha was accumulated for later dividends distribution assert_eq!(