use super::*;
use crate::nns::governance::NnsGovernanceCollectionStatus;
use crate::nns::neuron::{
NNS_NEURON_COLLECTION_STATE_SCHEMA_VERSION, NNS_NEURON_DISTRIBUTION_REPORT_SCHEMA_VERSION,
NnsNeuronCollectionState, NnsNeuronCollectionStep, NnsNeuronDistributionError,
NnsNeuronDistributionReport, advance_nns_neuron_collection_with_source,
build_nns_neuron_distribution_report,
};
struct PagedSource(bool);
impl NnsNeuronSource for PagedSource {
fn fetch_neuron_page<'a>(
&'a self,
request: &'a NnsGovernanceRequest,
start: Option<u64>,
size: u32,
) -> NnsNeuronSourceFuture<'a, NnsNeuronPage> {
Box::pin(async move {
assert_eq!(size, 2);
let (ids, next) = match start {
None => (vec![1, 2], Some(2)),
Some(2) => (vec![3], None),
other => panic!("unexpected cursor: {other:?}"),
};
let page = NnsNeuronPage {
neurons: ids
.into_iter()
.map(|id| NnsNeuronRow {
stake_e8s: 100_000_000,
staked_maturity_e8s_equivalent: None,
..sample_neuron(id)
})
.collect(),
next_start_neuron_id: next,
};
let mut data = match &request.source {
NnsGovernanceSourceSelection::ReplicaQuery { .. } => {
fixture_source_data(request, page)
}
NnsGovernanceSourceSelection::ReplicatedInterCanisterCall => {
NnsGovernanceSourceData::new(
page,
NnsGovernanceSourceProvenance::ReplicatedInterCanisterCall {
collector_canister_id: "aaaaa-aa".to_string(),
},
)
}
};
if self.0 && start.is_some() {
data.provenance = NnsGovernanceSourceProvenance::ReplicatedInterCanisterCall {
collector_canister_id: "2vxsx-fae".to_string(),
};
}
Ok(data)
})
}
fn fetch_neuron<'a>(
&'a self,
_: &'a NnsGovernanceRequest,
_: u64,
) -> NnsNeuronSourceFuture<'a, NnsNeuronRow> {
panic!("collection does not fetch detail")
}
}
#[tokio::test(flavor = "current_thread")]
async fn neuron_collection_state_resumes_until_api_exhaustion() {
let first_request = governance_request(1_700_000_000);
let state =
NnsNeuronCollectionState::new(&first_request, 2, 3).expect("valid neuron collection state");
let first =
advance_nns_neuron_collection_with_source(&first_request, &state, &PagedSource(false))
.await
.expect("first neuron collection page");
assert_eq!(
NNS_NEURON_COLLECTION_STATE_SCHEMA_VERSION,
first.state.schema_version()
);
assert_eq!(
first.state.status(),
NnsGovernanceCollectionStatus::Collecting
);
assert_eq!(first.state.pages_fetched(), 1);
let neurons_fetched: u64 = first.state.neurons_fetched();
assert_eq!(neurons_fetched, 2);
assert_eq!(first.state.next_start_neuron_id(), Some(2));
assert_eq!(first.state.network(), "ic");
assert_eq!(
first.state.governance_canister_id(),
"rrkah-fqaaa-aaaaa-aaaaq-cai"
);
assert_eq!(first.state.requested_source(), &first_request.source);
assert_eq!(
first.state.source(),
Some(&fixture_source_data(&first_request, ()).provenance)
);
assert_eq!(first.state.page_size(), 2);
assert!(!first.state.is_complete());
let persisted = serde_json::to_string(&first).expect("serialize collection step");
let persisted_json: serde_json::Value =
serde_json::from_str(&persisted).expect("parse serialized collection step");
assert_eq!(persisted_json["state"]["schema_version"], 1);
assert_eq!(
persisted_json["state"]["requested_source"]["source_transport"],
"replica_query"
);
let resumed: NnsNeuronCollectionStep =
serde_json::from_str(&persisted).expect("deserialize collection step");
let second_request = governance_request(1_700_000_100);
let second = advance_nns_neuron_collection_with_source(
&second_request,
&resumed.state,
&PagedSource(false),
)
.await
.expect("final neuron collection page");
assert_eq!(second.page.neurons.len(), 1);
assert_eq!(
second.state.status(),
NnsGovernanceCollectionStatus::Complete
);
assert_eq!(second.state.pages_fetched(), 2);
assert_eq!(second.state.neurons_fetched(), 3);
assert_eq!(second.state.next_start_neuron_id(), None);
assert_eq!(second.state.started_at(), "2023-11-14T22:13:20Z");
assert_eq!(second.state.updated_at(), "2023-11-14T22:15:00Z");
assert!(second.state.is_complete());
assert_neuron_distribution(&first.page.neurons, &second.page.neurons, &second.state);
let restored: NnsNeuronCollectionState = serde_json::from_value(
serde_json::to_value(&second.state).expect("serialize completed state"),
)
.expect("restore completed state");
let error =
advance_nns_neuron_collection_with_source(&second_request, &restored, &PagedSource(false))
.await
.expect_err("complete collection cannot make another source call");
assert!(matches!(
error,
NnsNeuronError::CollectionComplete { pages_fetched: 2 }
));
}
#[tokio::test(flavor = "current_thread")]
async fn neuron_collection_stops_at_explicit_page_limit() {
let request = governance_request(1_700_000_000);
assert!(matches!(
NnsNeuronCollectionState::new(&request, 2, 0),
Err(NnsNeuronError::InvalidCollectionMaxPages)
));
let state = NnsNeuronCollectionState::new(&request, 2, 1).expect("bounded collection state");
let step = advance_nns_neuron_collection_with_source(&request, &state, &PagedSource(false))
.await
.expect("bounded collection page");
assert_eq!(
step.state.status(),
NnsGovernanceCollectionStatus::PageLimitReached
);
assert_eq!(step.state.max_pages(), 1);
assert!(!step.state.is_complete());
let restored: NnsNeuronCollectionState = serde_json::from_value(
serde_json::to_value(&step.state).expect("serialize page-limited state"),
)
.expect("restore page-limited state");
let error = advance_nns_neuron_collection_with_source(&request, &restored, &PagedSource(false))
.await
.expect_err("page-limited collection cannot advance");
assert!(matches!(
error,
NnsNeuronError::CollectionPageLimitReached {
pages_fetched: 1,
max_pages: 1
}
));
let changed_request = NnsGovernanceRequest::replica_query_from_unix_secs(
"ic",
"https://example.com",
1_700_000_100,
"ic-query",
);
let error =
advance_nns_neuron_collection_with_source(&changed_request, &state, &PagedSource(false))
.await
.expect_err("collection identity must be fixed before source invocation");
assert!(matches!(
error,
NnsNeuronError::CollectionRequestMismatch {
field: "requested_source",
..
}
));
let mut invalid_json = serde_json::to_value(&step.state).expect("serialize collection state");
invalid_json["next_start_neuron_id"] = serde_json::Value::Null;
let invalid_state: NnsNeuronCollectionState =
serde_json::from_value(invalid_json).expect("deserialize modified collection state");
let error =
advance_nns_neuron_collection_with_source(&request, &invalid_state, &PagedSource(false))
.await
.expect_err("a full page cannot be changed into false completion");
assert!(matches!(
error,
NnsNeuronError::InvalidCollectionState { .. }
));
}
#[tokio::test(flavor = "current_thread")]
async fn neuron_collection_rejects_a_changed_canister_collector() {
let first_request =
NnsGovernanceRequest::replicated_inter_canister_call_from_unix_secs("ic", 1_700_000_000);
let state = NnsNeuronCollectionState::new(&first_request, 2, 3)
.expect("valid canister collection state");
let first =
advance_nns_neuron_collection_with_source(&first_request, &state, &PagedSource(true))
.await
.expect("first canister neuron page");
let second_request =
NnsGovernanceRequest::replicated_inter_canister_call_from_unix_secs("ic", 1_700_000_100);
let error = advance_nns_neuron_collection_with_source(
&second_request,
&first.state,
&PagedSource(true),
)
.await
.expect_err("collector principal must remain fixed");
assert!(matches!(
error,
NnsNeuronError::CollectionSourceChanged { .. }
));
}
fn assert_neuron_distribution(
first_page: &[NnsNeuronRow],
final_page: &[NnsNeuronRow],
state: &NnsNeuronCollectionState,
) {
let neurons = first_page
.iter()
.chain(final_page)
.cloned()
.collect::<Vec<_>>();
let distribution: NnsNeuronDistributionReport =
build_nns_neuron_distribution_report(state, &neurons)
.expect("portable public-neuron distribution");
assert_eq!(
distribution.context.schema_version,
NNS_NEURON_DISTRIBUTION_REPORT_SCHEMA_VERSION
);
assert_eq!(distribution.collected_neuron_count, 3);
assert_eq!(distribution.total_effective_stake_e8s, 300_000_000);
assert!(!distribution.context.point_in_time_guaranteed);
let json = serde_json::to_value(&distribution).expect("serialize neuron distribution");
assert_eq!(json["source"]["source_transport"], "replica_query");
assert!(matches!(
build_nns_neuron_distribution_report(state, &neurons[..2]),
Err(NnsNeuronDistributionError::NeuronCountMismatch { .. })
));
}