fix(iota-e2e-tests): fixed framework update tests (#3736)

* fix(iota-e2e-tests): fixed framework update tests data

* fix(iota-core): assert in case if an advance epoch transaction was executed in the safe mode

* fix(iota-core): code cleanup

crates/iota-core/src/authority.rs

@@ -4514,7 +4514,11 @@ impl AuthorityState {
         gas_cost_summary: &GasCostSummary,
         checkpoint: CheckpointSequenceNumber,
         epoch_start_timestamp_ms: CheckpointTimestamp,
-    ) -> anyhow::Result<(IotaSystemState, SystemEpochInfoEventV1, TransactionEffects)> {
+    ) -> anyhow::Result<(
+        IotaSystemState,
+        Option<SystemEpochInfoEventV1>,
+        TransactionEffects,
+    )> {
         let mut txns = Vec::new();
 
         if let Some(tx) = self.create_authenticator_state_tx(epoch_store) {
@@ -4644,11 +4648,13 @@ impl AuthorityState {
             .data
             .iter()
             .find(|event| event.is_system_epoch_info_event())
-            .expect("end of epoch tx must emit system epoch info event");
-        let system_epoch_info_event = bcs::from_bytes::<SystemEpochInfoEventV1>(
-            &system_epoch_info_event.contents,
-        )
-        .expect("deserialization should succeed since we asserted that the event is of this type");
+            .map(|event| {
+                bcs::from_bytes::<SystemEpochInfoEventV1>(&event.contents)
+                    .expect("event deserialization should succeed as type was pre-validated")
+            });
+        // The system epoch info event can be `None` in case if the `advance_epoch`
+        // Move function call failed and was executed in the safe mode.
+        assert!(system_epoch_info_event.is_some() || system_obj.safe_mode());
 
         // We must write tx and effects to the state sync tables so that state sync is
         // able to deliver to the transaction to CheckpointExecutor after it is

crates/iota-core/src/checkpoints/mod.rs

@@ -1445,11 +1445,16 @@ impl CheckpointBuilder {
                     )
                     .await?;
 
+                // The system epoch info event can be `None` in case if the `advance_epoch`
+                // Move function call failed and was executed in the safe mode.
+                // In this case, the tokens supply should be unchanged.
+                //
                 // SAFETY: The number of minted and burnt tokens easily fit into an i64 and due
                 // to those small numbers, no overflows will occur during conversion or
                 // subtraction.
-                let epoch_supply_change = system_epoch_info_event.minted_tokens_amount as i64
-                    - system_epoch_info_event.burnt_tokens_amount as i64;
+                let epoch_supply_change = system_epoch_info_event.map_or(0, |event| {
+                    event.minted_tokens_amount as i64 - event.burnt_tokens_amount as i64
+                });
 
                 let committee = system_state_obj
                     .get_current_epoch_committee()
@@ -1574,7 +1579,7 @@ impl CheckpointBuilder {
         checkpoint_effects: &mut Vec<TransactionEffects>,
         signatures: &mut Vec<Vec<GenericSignature>>,
         checkpoint: CheckpointSequenceNumber,
-    ) -> anyhow::Result<(IotaSystemState, SystemEpochInfoEventV1)> {
+    ) -> anyhow::Result<(IotaSystemState, Option<SystemEpochInfoEventV1>)> {
         let (system_state, system_epoch_info_event, effects) = self
             .state
             .create_and_execute_advance_epoch_tx(

crates/iota-e2e-tests/tests/framework_upgrades/mock_iota_systems/base/sources/genesis.move

@@ -3,12 +3,11 @@
 // SPDX-License-Identifier: Apache-2.0
 
 module iota_system::genesis {
-    use std::vector;
-    use iota::balance::{Self, Balance};
-    use iota::object::UID;
-    use iota::iota::IOTA;
-    use iota::tx_context::{Self, TxContext};
-    use std::option::Option;
+    use std::string::String;
+
+    use iota::balance;
+    use iota::iota::IotaTreasuryCap;
+    use iota::timelock::SystemTimelockCap;
 
     use iota_system::iota_system;
     use iota_system::validator;
@@ -55,15 +54,20 @@ module iota_system::genesis {
     public struct TokenAllocation has drop {
         recipient_address: address,
         amount_nanos: u64,
+
         staked_with_validator: Option<address>,
+        staked_with_timelock_expiration: Option<u64>,
     }
 
+    #[allow(unused_function)]
     fun create(
         iota_system_state_id: UID,
-        mut iota_supply: Balance<IOTA>,
+        mut iota_treasury_cap: IotaTreasuryCap,
         genesis_chain_parameters: GenesisChainParameters,
         genesis_validators: vector<GenesisValidatorMetadata>,
         _token_distribution_schedule: TokenDistributionSchedule,
+        _timelock_genesis_label: Option<String>,
+        system_timelock_cap: SystemTimelockCap,
         ctx: &mut TxContext,
     ) {
         assert!(tx_context::epoch(ctx) == 0, 0);
@@ -97,7 +101,7 @@ module iota_system::genesis {
                 network_address,
                 p2p_address,
                 primary_address,
-                balance::split(&mut iota_supply, 2500),
+                iota_treasury_cap.mint_balance(2500, ctx),
                 ctx
             );
 
@@ -108,11 +112,13 @@ module iota_system::genesis {
 
         iota_system::create(
             iota_system_state_id,
+            iota_treasury_cap,
             validators,
-            iota_supply, // storage_fund
+            balance::zero(),
             genesis_chain_parameters.protocol_version,
             genesis_chain_parameters.chain_start_timestamp_ms,
             genesis_chain_parameters.epoch_duration_ms,
+            system_timelock_cap,
             ctx,
         );
     }

crates/iota-e2e-tests/tests/framework_upgrades/mock_iota_systems/base/sources/iota_system.move

@@ -3,18 +3,16 @@
 // SPDX-License-Identifier: Apache-2.0
 
 module iota_system::iota_system {
-    use std::vector;
-
     use iota::balance::Balance;
-    use iota::object::UID;
-    use iota::iota::IOTA;
-    use iota::transfer;
-    use iota::tx_context::{Self, TxContext};
     use iota::dynamic_field;
+    use iota::iota::IOTA;
+    use iota::iota::IotaTreasuryCap;
+    use iota::timelock::SystemTimelockCap;
 
-    use iota_system::validator::Validator;
-    use iota_system::iota_system_state_inner::IotaSystemStateInner;
-    use iota_system::iota_system_state_inner;
+    use iota_system::validator::ValidatorV1;
+    use iota_system::iota_system_state_inner::{Self, IotaSystemStateV1};
+
+    const SYSTEM_TIMELOCK_CAP_DF_KEY: vector<u8> = b"sys_timelock_cap";
 
     public struct IotaSystemState has key {
         id: UID,
@@ -23,14 +21,17 @@ module iota_system::iota_system {
 
     public(package) fun create(
         id: UID,
-        validators: vector<Validator>,
+        iota_treasury_cap: IotaTreasuryCap,
+        validators: vector<ValidatorV1>,
         storage_fund: Balance<IOTA>,
         protocol_version: u64,
         epoch_start_timestamp_ms: u64,
         epoch_duration_ms: u64,
+        system_timelock_cap: SystemTimelockCap,
         ctx: &mut TxContext,
     ) {
         let system_state = iota_system_state_inner::create(
+            iota_treasury_cap,
             validators,
             storage_fund,
             protocol_version,
@@ -44,21 +45,22 @@ module iota_system::iota_system {
             version,
         };
         dynamic_field::add(&mut self.id, version, system_state);
+        dynamic_field::add(&mut self.id, SYSTEM_TIMELOCK_CAP_DF_KEY, system_timelock_cap);
         transfer::share_object(self);
     }
 
+    #[allow(unused_function)]
     fun advance_epoch(
+        validator_target_reward: u64,
         storage_charge: Balance<IOTA>,
         computation_reward: Balance<IOTA>,
         wrapper: &mut IotaSystemState,
         new_epoch: u64,
         next_protocol_version: u64,
         storage_rebate: u64,
-        _non_refundable_storage_fee: u64,
-        _storage_fund_reinvest_rate: u64, // share of storage fund's rewards that's reinvested
-                                         // into storage fund, in basis point.
-        _reward_slashing_rate: u64, // how much rewards are slashed to punish a validator, in bps.
-        epoch_start_timestamp_ms: u64, // Timestamp of the epoch start
+        non_refundable_storage_fee: u64,
+        reward_slashing_rate: u64,
+        epoch_start_timestamp_ms: u64,
         ctx: &mut TxContext,
     ) : Balance<IOTA> {
         let self = load_system_state_mut(wrapper);
@@ -67,28 +69,32 @@ module iota_system::iota_system {
             self,
             new_epoch,
             next_protocol_version,
+            validator_target_reward,
             storage_charge,
             computation_reward,
             storage_rebate,
+            non_refundable_storage_fee,
+            reward_slashing_rate,
             epoch_start_timestamp_ms,
+            ctx
         );
 
         storage_rebate
     }
 
-    public fun active_validator_addresses(wrapper: &mut IotaSystemState): vector<address> {
+    public fun active_validator_addresses(_wrapper: &mut IotaSystemState): vector<address> {
         vector::empty()
     }
 
-    fun load_system_state_mut(self: &mut IotaSystemState): &mut IotaSystemStateInner {
+    fun load_system_state_mut(self: &mut IotaSystemState): &mut IotaSystemStateV1 {
         load_inner_maybe_upgrade(self)
     }
 
-    fun load_inner_maybe_upgrade(self: &mut IotaSystemState): &mut IotaSystemStateInner {
+    fun load_inner_maybe_upgrade(self: &mut IotaSystemState): &mut IotaSystemStateV1 {
         let version = self.version;
         // TODO: This is where we check the version and perform upgrade if necessary.
 
-        let inner: &mut IotaSystemStateInner = dynamic_field::borrow_mut(&mut self.id, version);
+        let inner: &mut IotaSystemStateV1 = dynamic_field::borrow_mut(&mut self.id, version);
         assert!(iota_system_state_inner::system_state_version(inner) == version, 0);
         inner
     }