use crate::{
cdk::candid::Principal,
cdk::types::Cycles,
config::schema::{
CanisterAuthConfig, CanisterConfig, CanisterKind, CyclesFundingPolicyConfig,
DiagnosticsCanisterConfig, MetricsCanisterConfig, ShardPool, ShardPoolPolicy,
ShardingConfig, StandardsCanisterConfig,
},
ids::{CanisterRole, CanonicalNetworkId, FleetId, FleetKey, SubnetSlotId},
ops::runtime::env::EnvOps,
storage::stable::env::{EnvData, EnvRecord},
test::config::ConfigTestBuilder,
};
#[must_use]
pub fn fleet_key(byte: u8) -> FleetKey {
FleetKey {
network: CanonicalNetworkId::public_ic(),
fleet_id: FleetId::from_generated_bytes([byte.saturating_add(1); 32]),
}
}
pub fn init_sharding_test_config() {
let mut sharding = ShardingConfig::default();
sharding.pools.insert(
"primary".to_string(),
ShardPool {
canister_role: CanisterRole::from("shard"),
policy: ShardPoolPolicy {
capacity: 1,
initial_shards: 1,
max_shards: 2,
},
},
);
let root_cfg = CanisterConfig {
kind: CanisterKind::Root,
initial_cycles: Cycles::new(5_000_000_000_000),
topup: None,
icp_refill: None,
cycles_funding: CyclesFundingPolicyConfig::default(),
scaling: None,
sharding: None,
binding: None,
auth: CanisterAuthConfig::default(),
standards: StandardsCanisterConfig::default(),
diagnostics: DiagnosticsCanisterConfig::default(),
metrics: MetricsCanisterConfig::default(),
};
let manager_cfg = CanisterConfig {
kind: CanisterKind::Service,
initial_cycles: Cycles::new(5_000_000_000_000),
topup: None,
icp_refill: None,
cycles_funding: CyclesFundingPolicyConfig::default(),
scaling: None,
sharding: Some(sharding),
binding: None,
auth: CanisterAuthConfig::default(),
standards: StandardsCanisterConfig::default(),
diagnostics: DiagnosticsCanisterConfig::default(),
metrics: MetricsCanisterConfig::default(),
};
let shard_cfg = CanisterConfig {
kind: CanisterKind::Shard,
initial_cycles: Cycles::new(5_000_000_000_000),
topup: None,
icp_refill: None,
cycles_funding: CyclesFundingPolicyConfig::default(),
scaling: None,
sharding: None,
binding: None,
auth: CanisterAuthConfig::default(),
standards: StandardsCanisterConfig::default(),
diagnostics: DiagnosticsCanisterConfig::default(),
metrics: MetricsCanisterConfig::default(),
};
let _config = ConfigTestBuilder::new()
.with_default_canister(CanisterRole::ROOT, root_cfg)
.with_default_canister("manager", manager_cfg)
.with_default_canister("shard", shard_cfg)
.install();
let root_pid = Principal::from_slice(&[1; 29]);
import_test_env("manager", SubnetSlotId::DEFAULT, root_pid);
}
pub fn import_test_env(
canister_role: impl Into<CanisterRole>,
subnet_slot: impl Into<SubnetSlotId>,
root_pid: Principal,
) {
let snapshot = EnvRecord {
canister_role: Some(canister_role.into()),
subnet_slot: Some(subnet_slot.into()),
root_pid: Some(root_pid),
fleet_root_pid: Some(root_pid),
subnet_pid: Some(root_pid),
parent_pid: Some(root_pid),
};
EnvOps::import(EnvData { record: snapshot }).expect("import test env");
}