governance needs to deserialize to its proto type, not the api type

dfinity · Jan 10, 2025 · ec8228c · ec8228c
1 parent 4b9b18f
commit ec8228c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 17 deletions.
diff --git a/rs/sns/governance/canister/canister.rs b/rs/sns/governance/canister/canister.rs
@@ -35,10 +35,9 @@ use ic_sns_governance_api::pb::v1::{
     GetModeResponse, GetNeuron, GetNeuronResponse, GetProposal, GetProposalResponse,
     GetRunningSnsVersionRequest, GetRunningSnsVersionResponse,
     GetSnsInitializationParametersRequest, GetSnsInitializationParametersResponse,
-    GetUpgradeJournalRequest, GetUpgradeJournalResponse, Governance as GovernanceProto,
-    ListNervousSystemFunctionsResponse, ListNeurons, ListNeuronsResponse, ListProposals,
-    ListProposalsResponse, ManageNeuron, ManageNeuronResponse, NervousSystemParameters,
-    RewardEvent, SetMode, SetModeResponse,
+    GetUpgradeJournalRequest, GetUpgradeJournalResponse, ListNervousSystemFunctionsResponse,
+    ListNeurons, ListNeuronsResponse, ListProposals, ListProposalsResponse, ManageNeuron,
+    ManageNeuronResponse, NervousSystemParameters, RewardEvent, SetMode, SetModeResponse,
 };
 #[cfg(feature = "test")]
 use ic_sns_governance_api::pb::v1::{
@@ -211,8 +210,7 @@ fn caller() -> PrincipalId {
 /// In contrast to canister_init(), this method does not do deserialization.
 /// In addition to canister_init, this method is called by canister_post_upgrade.
 #[init]
-fn canister_init_(init_payload: GovernanceProto) {
-    let init_payload = sns_gov_pb::Governance::from(init_payload);
+fn canister_init_(init_payload: sns_gov_pb::Governance) {
     let init_payload = ValidGovernanceProto::try_from(init_payload).expect(
         "Cannot start canister, because the deserialized \
          GovernanceProto is invalid in some way",
@@ -277,7 +275,7 @@ fn canister_post_upgrade() {
 
     let reader = BufferedStableMemReader::new(STABLE_MEM_BUFFER_SIZE);
 
-    match GovernanceProto::decode(reader) {
+    match sns_gov_pb::Governance::decode(reader) {
         Err(err) => {
             log!(
                 ERROR,
@@ -304,7 +302,7 @@ fn canister_post_upgrade() {
     log!(INFO, "Completed post upgrade");
 }
 
-fn populate_finalize_disbursement_timestamp_seconds(governance_proto: &mut GovernanceProto) {
+fn populate_finalize_disbursement_timestamp_seconds(governance_proto: &mut sns_gov_pb::Governance) {
     for neuron in governance_proto.neurons.values_mut() {
         for disbursement in neuron.disburse_maturity_in_progress.iter_mut() {
             disbursement.finalize_disbursement_timestamp_seconds = Some(

diff --git a/rs/sns/governance/canister/tests.rs b/rs/sns/governance/canister/tests.rs
@@ -1,7 +1,11 @@
 use super::*;
 use assert_matches::assert_matches;
 use candid_parser::utils::{service_equal, CandidSource};
-use ic_sns_governance_api::pb::v1::{DisburseMaturityInProgress, Neuron};
+use ic_sns_governance::pb::v1::{
+    governance::{Version, Versions},
+    upgrade_journal_entry::{Event, UpgradeStepsRefreshed},
+    DisburseMaturityInProgress, Neuron, UpgradeJournal, UpgradeJournalEntry,
+};
 use maplit::btreemap;
 use pretty_assertions::assert_eq;
 use std::collections::HashSet;
@@ -61,7 +65,7 @@ fn test_set_time_warp() {
 fn test_populate_finalize_disbursement_timestamp_seconds() {
     // Step 1: prepare a neuron with 2 in progress disbursement, one with
     // finalize_disbursement_timestamp_seconds as None, and the other has incorrect timestamp.
-    let mut governance_proto = GovernanceProto {
+    let mut governance_proto = sns_gov_pb::Governance {
         neurons: btreemap! {
             "1".to_string() => Neuron {
                 disburse_maturity_in_progress: vec![
@@ -86,7 +90,7 @@ fn test_populate_finalize_disbursement_timestamp_seconds() {
     populate_finalize_disbursement_timestamp_seconds(&mut governance_proto);
 
     // Step 3: verifies that both disbursements have the correct finalization timestamps.
-    let expected_governance_proto = GovernanceProto {
+    let expected_governance_proto = sns_gov_pb::Governance {
         neurons: btreemap! {
             "1".to_string() => Neuron {
                 disburse_maturity_in_progress: vec![
@@ -111,12 +115,6 @@ fn test_populate_finalize_disbursement_timestamp_seconds() {
 
 #[test]
 fn test_upgrade_journal() {
-    use ic_sns_governance::pb::v1::{
-        governance::{Version, Versions},
-        upgrade_journal_entry::{Event, UpgradeStepsRefreshed},
-        UpgradeJournal, UpgradeJournalEntry,
-    };
-
     let journal = UpgradeJournal {
         entries: vec![UpgradeJournalEntry {
             timestamp_seconds: Some(1000),