diff --git a/rs/sns/governance/canister/canister.rs b/rs/sns/governance/canister/canister.rs index 2c5e83ebaf4..3a0b80b66ad 100644 --- 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 index f5c7e7a90c2..496348898b1 100644 --- 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),