From 8f779134294764b298251e2e7159b176a6dc8cf5 Mon Sep 17 00:00:00 2001 From: Max Summe Date: Mon, 3 Mar 2025 15:30:58 -0800 Subject: [PATCH] Tests working (mostly) --- .../api/src/ic_nns_governance.pb.v1.rs | 5 --- .../src/gen/ic_nns_governance.pb.v1.rs | 3 ++ rs/nns/governance/src/pb/conversions.rs | 1 - rs/nns/governance/src/reward/distribution.rs | 2 + .../src/timer_tasks/distribute_rewards.rs | 43 +++++++++++++++++++ rs/nns/governance/tests/governance.rs | 7 ++- 6 files changed, 51 insertions(+), 10 deletions(-) diff --git a/rs/nns/governance/api/src/ic_nns_governance.pb.v1.rs b/rs/nns/governance/api/src/ic_nns_governance.pb.v1.rs index 4e4f50c13f6..ba8dbc41097 100644 --- a/rs/nns/governance/api/src/ic_nns_governance.pb.v1.rs +++ b/rs/nns/governance/api/src/ic_nns_governance.pb.v1.rs @@ -2458,11 +2458,6 @@ pub struct RewardEvent { /// In both of these cases, the rewards purse rolls over into the next round. #[prost(uint64, optional, tag = "6")] pub rounds_since_last_distribution: Option, - - /// Whether or not the asynchronous distribution of rewards has completed. Rewards are - /// calculated once a day, but are distributed after the calculation in batches. - #[prost(bool, optional, tag = "8")] - pub distribution_complete: Option, } #[derive(candid::CandidType, candid::Deserialize, serde::Serialize, comparable::Comparable)] #[allow(clippy::derive_partial_eq_without_eq)] diff --git a/rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs b/rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs index b6865050d21..f442dc3d22f 100644 --- a/rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs +++ b/rs/nns/governance/src/gen/ic_nns_governance.pb.v1.rs @@ -4273,6 +4273,9 @@ pub struct Account { #[prost(message, optional, tag = "2")] pub subaccount: ::core::option::Option, } +/// A reward disbribution that has been calculated but not fully disbursed. +/// This supports large reward distributions that may need to be split into multiple +/// messages. #[derive( candid::CandidType, candid::Deserialize, diff --git a/rs/nns/governance/src/pb/conversions.rs b/rs/nns/governance/src/pb/conversions.rs index 21618cf340f..c179d7a2106 100644 --- a/rs/nns/governance/src/pb/conversions.rs +++ b/rs/nns/governance/src/pb/conversions.rs @@ -1985,7 +1985,6 @@ impl From for pb_api::RewardEvent { total_available_e8s_equivalent: item.total_available_e8s_equivalent, latest_round_available_e8s_equivalent: item.latest_round_available_e8s_equivalent, rounds_since_last_distribution: item.rounds_since_last_distribution, - distribution_complete: None, } } } diff --git a/rs/nns/governance/src/reward/distribution.rs b/rs/nns/governance/src/reward/distribution.rs index 08c85b11426..b6dcf4dc547 100644 --- a/rs/nns/governance/src/reward/distribution.rs +++ b/rs/nns/governance/src/reward/distribution.rs @@ -2,6 +2,7 @@ use crate::governance::{Governance, LOG_PREFIX}; use crate::neuron_store::NeuronStore; use crate::pb::v1::RewardsDistributionInProgress; use crate::storage::with_rewards_distribution_state_machine_mut; +#[cfg(not(feature = "canbench-rs"))] use crate::timer_tasks::run_distribute_rewards_periodic_task; use ic_nns_common::pb::v1::NeuronId; use ic_stable_structures::storable::Bound; @@ -29,6 +30,7 @@ impl Governance { println!("{}Error scheduling rewards distribution: {}", LOG_PREFIX, e); } + #[cfg(not(feature = "canbench-rs"))] run_distribute_rewards_periodic_task(); } diff --git a/rs/nns/governance/src/timer_tasks/distribute_rewards.rs b/rs/nns/governance/src/timer_tasks/distribute_rewards.rs index 4bd75e74f1c..b57dbc9ec47 100644 --- a/rs/nns/governance/src/timer_tasks/distribute_rewards.rs +++ b/rs/nns/governance/src/timer_tasks/distribute_rewards.rs @@ -54,3 +54,46 @@ impl PeriodicSyncTask for DistributeRewardsTask { const NAME: &'static str = "distribute_rewards"; const INTERVAL: Duration = Duration::from_secs(2); } + +#[cfg(test)] +mod tests { + use super::*; + use crate::canister_state::{governance_mut, set_governance_for_tests}; + use crate::governance::Governance; + use crate::reward::distribution::RewardsDistribution; + use crate::test_utils::{MockEnvironment, MockRandomness, StubCMC, StubIcpLedger}; + use crate::timer_tasks::distribute_rewards::REWARDS_TIMER_ID; + use ic_nervous_system_timers::test::run_pending_timers_every_interval_for_count; + use ic_nns_common::pb::v1::NeuronId; + use std::sync::Arc; + + #[test] + fn test_reward_scheduling_and_cancelling() { + let governance_proto = crate::pb::v1::Governance::default(); + + let governance = Governance::new( + governance_proto, + Arc::new(MockEnvironment::new(Default::default(), 0)), + Arc::new(StubIcpLedger {}), + Arc::new(StubCMC {}), + Box::new(MockRandomness::new()), + ); + + set_governance_for_tests(governance); + let governance = governance_mut(); + + // In this test, we don't care that rewards are actually distributed, only that the + // timer is scheduled and then cancelled. Other tests cover that the rewards are distributed. + let mut distribution = RewardsDistribution::new(); + for id in 0..10 { + distribution.add_reward(NeuronId { id }, 10); + } + // create 2 distributions + governance.schedule_pending_rewards_distribution(1, distribution.clone()); + assert!(REWARDS_TIMER_ID.with(|id| id.borrow().is_some())); + + run_pending_timers_every_interval_for_count(DistributeRewardsTask::INTERVAL, 3); + + assert!(REWARDS_TIMER_ID.with(|id| id.borrow().is_none())); + } +} diff --git a/rs/nns/governance/tests/governance.rs b/rs/nns/governance/tests/governance.rs index e72d1ee7ed4..402df9a3221 100644 --- a/rs/nns/governance/tests/governance.rs +++ b/rs/nns/governance/tests/governance.rs @@ -3647,9 +3647,8 @@ async fn test_reward_event_proposals_last_longer_than_reward_period() { gov.run_periodic_tasks().now_or_never(); assert_eq!(gov.latest_reward_event().day_after_genesis, 6); // let's advance far enough to trigger a reward event - fake_driver.advance_time_by(REWARD_DISTRIBUTION_PERIOD_SECONDS); - run_pending_timers(); - gov.run_periodic_tasks().now_or_never(); + fake_driver.advance_time_by(REWARD_DISTRIBUTION_PERIOD_SECONDS + 1); + run_pending_timers_every_interval_for_count(std::time::Duration::from_secs(3), 3); // Inspect latest_reward_event. let fully_elapsed_reward_rounds = 7; @@ -3708,7 +3707,7 @@ async fn test_reward_event_proposals_last_longer_than_reward_period() { // Now let's advance again -- a new empty reward event should happen fake_driver.advance_time_by(REWARD_DISTRIBUTION_PERIOD_SECONDS); run_pending_timers(); - gov.run_periodic_tasks().now_or_never(); + assert_eq!( *gov.latest_reward_event(), RewardEvent {