Tests working (mostly)

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<u64>,
-
-    /// 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<bool>,
 }
 #[derive(candid::CandidType, candid::Deserialize, serde::Serialize, comparable::Comparable)]
 #[allow(clippy::derive_partial_eq_without_eq)]

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<Subaccount>,
 }
+/// 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,

rs/nns/governance/src/pb/conversions.rs

@@ -1985,7 +1985,6 @@ impl From<pb::RewardEvent> 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,
         }
     }
 }

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();
     }
 

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()));
+    }
+}

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 {